`clip` throws due to changes in `jax`?

[dom-linkevicius]

Hi, I was trying to run the first basic tutorial at the top (https://jaxley.readthedocs.io/en/latest/tutorials/01_morph_neurons.html), but

>>> v = jx.integrate(cell, delta_t=0.025)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/integrate.py", line 458, in integrate
    all_states, recordings = nested_checkpoint_scan(
                             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/utils/jax_utils.py", line 61, in nested_checkpoint_scan
    return _inner_nested_scan(f, init, sub_xs, nested_lengths, scan_fn, checkpoint_fn)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/utils/jax_utils.py", line 67, in _inner_nested_scan
    return scan_fn(f, init, xs, lengths[0])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/integrate.py", line 414, in _body_fun
    state = step_fn(state, all_params, externals, external_inds, delta_t)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/integrate.py", line 188, in step_fn
    state = module.step(
            ^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/modules/base.py", line 78, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/modules/base.py", line 2857, in step
    u, (linear_terms, const_terms) = self._step_channels(
                                     ^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/modules/base.py", line 3091, in _step_channels
    states = self._step_channels_state(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/modules/base.py", line 3131, in _step_channels_state
    states_updated = channel.update_states(
                     ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/channels/pospischil.py", line 96, in update_states
    new_m = solve_gate_exponential(m, dt, *self.m_gate(v, params["vt"]))
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/channels/pospischil.py", line 123, in m_gate
    alpha = 0.32 * efun(-0.25 * v_alpha) / 0.25
                   ^^^^^^^^^^^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/channels/pospischil.py", line 32, in efun
    return x / (save_exp(x) - 1.0)
                ^^^^^^^^^^^
  File "/home/dom/jaxley/lib/python3.12/site-packages/jaxley/solver_gate.py", line 9, in save_exp
    x = jnp.clip(x, a_max=max_value)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: clip() got an unexpected keyword argument 'a_max'

Looking at the jax code it seems like they have changed the kwarg names in clip. I'm running things on Python 3.12.3, jax 0.10.0 and Jaxley 0.13.0.

[jnsbck]

Hey, thanks for opening this.
We had some issues with jax compatibility in the past see #571. I am not sure if this has been resolved yet, @michaeldeistler do you know? For the moment I would suggest downgrading jax.

We should also make sure to update all the dependencies at some point as well and make sure tests still pass @michaeldeistler.

[dom-linkevicius]

@jnsbck Many thanks for your reply. Do you know what is the last combination of versions that is compatible?

[michaeldeistler]

I don't think this has been resolved, the 1.0.0 branch also still uses:

def save_exp(x, max_value: float = 20.0):
    """Clip the input to a maximum value and return its exponential."""
    x = jnp.clip(x, a_max=max_value)
    return jnp.exp(x)

And no, I don't know which version still works. However, I think this JAX change is quite recent.

[jnsbck]

According to the jax changelog a_max was deprecated in v0.10.0, so if you want to run on a recent version, v0.9.2 should do it. Otherwise consider pinning jax to a version <0.6.0 for the reasons mentioned above.