JAX 0.5.0: unsupported operand type(s) for *: 'tuple' and 'DynamicJaxprTracer'

[jecampagne]

Hello,
Here is a toy snippet that make me puzzled (jax 0.5.0):

import jax
import jax.numpy as jnp

h0=0.01

def body(val):
  t,y,sigma = val
  jax.debug.print("t {}",t)
  h = h0 * (t/(1.0 + (h0*(t-1.0))) )
  y *= (1.0+h)
  sigma -= 0.1
  t += 1.0
  return t,y,sigma

def cond(val):
  t,y,sigma = val
  return (sigma>0.1) & (sigma<2.0)

val=(1.,jnp.ones((2,2)),1.0)
jax.lax.while_loop(cond,body,val)

Here is what I get

Now, if I perform independantly of the while-loop use case the computation of "h" by a jitted function

@jax.jit
def f(t, h0):
  print(t,h0)
  return h0 * (t/(1.0 + (h0*(t-1.0))) )

f(0.1,0.1)

No crash occurs and the print gives indication of the use of DynamicJaxprTrace

Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace> Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace>

Array(0.01098901, dtype=float32, weak_type=True)

So, it seems that the line t,y,sigma = valof the body function turns out that "t" is considered as a Tuple?

This is bizarre because I used to do such things before 0.5.0, am I wrong? how I can do such simple computation in the while-loop? Thanks

[dfm]

I can't reproduce this error on any versions of JAX that I've tried (including 0.5.0 and 0.5.1) on colab or on my mac. I think that what you've written should work, and it does for me!

[jecampagne]

Thanks @dfm for your interest.
I will asap try to see if I have a side effect.
Now btw how do you have install a JAX version 0.5.x on Colab?

[jecampagne]

oh! I find my mistake :)
I had stupidly a comma h0=0.01, in my code, that I have not included in the snippet .... sorry.
Now the question on JAX 0.5.x version on Colab would be helpful.

[dfm]

Sure! To update JAX on colab, use:

%pip install -U jax

or, for GPU,

%pip install -U jax[cuda12]

at the top of your notebook.