GenCast autoregressive rollout fails with JaxArrayWrapper / DynamicJaxprTracer __array_ufunc__ TypeError on both GPU and Colab TPU

[DarinaAndr]

Hi, would really appreciate some help with the following issue of getting a TypeError while running the GenCast autoregressive rollout. It fails with a TypeError related to xarray_jax.JaxArrayWrapper interacting with a DynamicJaxprTracer.

The failure occurs during the autoregressive rollout step (chunked_prediction_generator_multiple_runs) in the GenCast demo.

I can reproduce the same error on a local multi-GPU machine, and on Google Colab TPU (v5e-1 TPU with 2025.07 runtime version), and I think the error is related to the GenCast rollout path interacting with xarray_jax.

The simpler GraphCast forward pass works correctly.

Execution fails with:

TypeError: operand type(s) all returned NotImplemented from __array_ufunc__(
<ufunc 'multiply'>, '__call__',
JitTracer(float32[]),
xarray_jax.JaxArrayWrapper(JitTracer(float32[1,1,181,360]))
): 'DynamicJaxprTracer', 'JaxArrayWrapper'

For reproduction:
Using GenCast demo notebook, the failure occurs in the Autoregressive rollout cell:

chunks = []
for chunk in rollout.chunked_prediction_generator_multiple_runs(
    predictor_fn=run_forward_pmap,
    rngs=rngs,
    inputs=eval_inputs,
    targets_template=eval_targets * np.nan,
    forcings=eval_forcings,
    num_steps_per_chunk=1,
    num_samples=num_ensemble_members,
    pmap_devices=jax.local_devices(),
):
    chunks.append(chunk)

predictions = xarray.combine_by_coords(chunks)

Setup:

Machine:

Linux cluster

2× NVIDIA GPUs
Python environment:

Python 3.11
jax 0.4.38
jaxlib 0.4.38
numpy 2.4.2
xarray 2024.11.0
pandas 2.2.3

Devices detected:

devices: [CudaDevice(id=0), CudaDevice(id=1)]
backend: gpu

Error traceback ends with:

TypeError: operand type(s) all returned NotImplemented from __array_ufunc__(
<ufunc 'multiply'>, '__call__',
JitTracer(float32[]),
xarray_jax.JaxArrayWrapper(JitTracer(float32[1,1,181,360]))
)

[mjwillson]

Hello! I was unable to reproduce this in a public v5e-1 colab TPU kernel, v5e-1.

After running the pip install cells in the colab, pip list shows

jax 0.9.1
jaxlib 0.9.1
numpy 2.0.2
xarray 2026.2.0
pandas 2.2.2

Looks like you're running a significantly older version of JAX, which may have some compatibility issues with xarray_jax -- would suggest upgrading it and seeing if that fixes things?

[DarinaAndr]

Hello! I was unable to reproduce this in a public v5e-1 colab TPU kernel, v5e-1.

After running the pip install cells in the colab, pip list shows

jax 0.9.1 jaxlib 0.9.1 numpy 2.0.2 xarray 2026.2.0 pandas 2.2.2

Looks like you're running a significantly older version of JAX, which may have some compatibility issues with xarray_jax -- would suggest upgrading it and seeing if that fixes things?

Hi @mjwillson, thank you for your reply! I tried downgrading to an older version of JAX to see if it would help while troubleshooting, but I have now rerun the code with pip list showing:

jax 0.9.1
jaxlib 0.9.1
numpy 2.0.2
xarray 2026.2.0
pandas 2.2.2
xarray-tensorstore 0.3.0

to match the working Colab stack (the latest runtime version)

However, I still get the exact same error on my GPU setup. I also verified that the failure reproduces in the older Colab runtime with the same traceback (in dpm_solver_plus_plus_2s.py during noise_levels[0] * spherical_white_noise_like(...), involving DynamicJaxprTracer and JaxArrayWrapper).

To reproduce the error in Colab, v5-e1TPU and Runtime version 2025/07 must be selected. The latest Colab runtime version works for me, but unfortunately that doesn't help with setting it up on my GPUs.

So it doesn’t seem to be explained by outdated JAX alone. Let me know if there’s another fix anyone can think of....