[JAX] Fix init2winit code for an upcoming jax.numpy.array() change.
[JAX] Fix init2winit code for an upcoming jax.numpy.array() change.

An upcoming change to jax.numpy.array() means that, under a transformation like jax.jit(), it will always stage its arrays into the trace. This often breaks if the array is being used for a shape calculation. In this case, we can just use a tuple.

PiperOrigin-RevId: 398004781

google/init2winit