Explicitly convert jax.numpy.meshgrid outputs to list

The return type of several jax.numpy APIs will change from list to tuple in an upcoming JAX version, following a similar change in NumPy 2.0.

PiperOrigin-RevId: 602334680

dm_pix/_src/augment.py

@@ -231,8 +231,9 @@ def elastic_deformation(
       sigma=sigma,
       kernel_size=kernel_size) * alpha
 
-  meshgrid = jnp.meshgrid(*[jnp.arange(size) for size in single_channel_shape],
-                          indexing="ij")
+  meshgrid = list(
+      jnp.meshgrid(
+          *[jnp.arange(size) for size in single_channel_shape], indexing="ij"))
   meshgrid[0] += shift_map_i
   meshgrid[1] += shift_map_j