removing jax.Array type to avoid the need of using the new jax version

JakobRobnik · Nov 27, 2023 · f300409 · f300409
1 parent 6fdbd35
commit f300409
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/mclmc/dynamics.py b/mclmc/dynamics.py
@@ -8,10 +8,10 @@
 class State(NamedTuple):
  """Dynamical state"""
 
- x: jax.Array
- u: jax.Array
+ x: any#jax.Array
+ u: any#jax.Array
  l: float
- g: jax.Array
+ g: any#jax.Array
  key: tuple
 
 

diff --git a/mclmc/tune.py b/mclmc/tune.py
@@ -12,7 +12,7 @@ class Hyperparameters(NamedTuple):
 
  L: float
  eps: float
- sigma: jax.Array
+ sigma: any
 
 
 # all tuning functions are wrappers, recieving some parameters and returning a function