Closed jeremiecoullon closed 3 years ago
Good catch! There's another incompatibility with newer versions of JAX (#36), so I pinned the version temporarily. That's the price to pay when you use a young library :)
Good news is these are easily fixable and won't impact performance 😅
Jax 0.2.6 stops supporting non-hashable static arguments for jitted functions. See changelog. That change log points to an explanation here.
This means that arrays cannot be set as static arguments in jitted functions. An example of a functions that does this at the moment: in
inference.warmup.num_steps_adaptation.py
, the functionrun()
defined in thelongest_batch_before_turn()
function sets arrays as static.This means that cloning mcx and running it with an up-to-date version of jax doesn't work.