Fix GPU memory leak that came up in RCN example

[StannisZhou]

The underlying issue seems to be known yet unresolved (jax-ml/jax#282).

Disabling jit and being careful about using jax arrays outside the inference function get rid of the leak. Inference is becoming a bit slower (sometimes as much as 2x slower) due to disabling jit for run_bp.

Also fixed a mypy issue (relaxed a Hashable to Any) and reduced precommit autoupdate frequence.

[codecov-commenter]

Codecov Report

Merging #97 (dec9c70) into master (2f389ab) will not change coverage.
The diff coverage is 100.00%.

@@            Coverage Diff            @@
##            master       #97   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            8         8           
  Lines          631       633    +2     
=========================================
+ Hits           631       633    +2     

Impacted Files	Coverage Δ
pgmax/fg/graph.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2f389ab...dec9c70. Read the comment docs.

[shrinuKushagra]

Looks good to me.
Seems like using numpy arrays instead of jax arrays in init functions fixed the error.

[shrinuKushagra]

Not expert on jax but for my own clarification.
Do we not need to clear the variables self.value and data from jax memory memory?

[StannisZhou]

self.value and data would stay on CPU. Intermediate objects on GPU will be cleared automatically if you use the platform option.

The leak came from a failure to clear large jit compilation cache (the cache is large because it involves a large constant array, i.e. the wiring). A completely leak-free way is to include wiring as part of the arguments, but in my experiments that leads to slow compiling and hurts performance. Current solution is a bit of a compromise but should suffice for most cases.