Error/Difference in semantics in jax.shard_map

[wittenator]

This ticket patrick-kidger/diffrax#735 shows an VMA error in a jax.shard_map with a jax.scan call that can be mitigated by setting check_vma=False. Interestingly, while passing the mesh to the shard_map and setting check_vma=False works for the example, automatically inferring the mesh and unsetting the VMA check fails with the following error:

import functools as ft

import jax
import jax.numpy as jnp
from diffrax import ODETerm, Tsit5, diffeqsolve
from jax import shard_map, sharding
from jax.experimental import mesh_utils

mesh = sharding.Mesh(mesh_utils.create_device_mesh((2,)), ["i"])
spec = sharding.PartitionSpec("i")

with jax.set_mesh(mesh):

    @jax.jit
    @ft.partial(shard_map, mesh=mesh, in_specs=spec, out_specs=spec, check_vma=False)
    @jax.vmap
    def run(y0):
        term = ODETerm(lambda t, y, args: -y)
        solver = Tsit5()
        t0 = 0
        t1 = 1
        dt0 = 0.1
        sol = diffeqsolve(term, solver, t0, t1, dt0, y0)
        return sol.ys

    y0 = jnp.array([10.0, 10.0])
    y0 = jax.device_put(y0, spec)
    out = run(y0)
    print(out)
    jax.debug.visualize_array_sharding(out)

As pointed out in patrick-kidger/diffrax#738 this may actually be a bug in Jax, but since I cannot reproduce it yet without diffrax I would like to open a discussion first. The stack trace looks as follows:

Traceback (most recent call last):
  File "/workspaces/gfdm-jax/scripts/test.py", line 28, in <module>
    out = run(y0)
          ^^^^^^^
  File "/workspaces/gfdm-jax/scripts/test.py", line 23, in run
    sol = diffeqsolve(term, solver, t0, t1, dt0, y0)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_integrate.py", line 1456, in diffeqsolve
    final_state, aux_stats = adjoint.loop(
                             ^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/equinox/_module/_prebuilt.py", line 34, in __call__
    return self.__func__(self.__self__, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_adjoint.py", line 299, in loop
    final_state = self._loop(
                  ^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_integrate.py", line 666, in loop
    _, traced_jump, traced_result = eqx.filter_eval_shape(body_fun_aux, init_state)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_integrate.py", line 373, in body_fun_aux
    (y, y_error, dense_info, solver_state, solver_result) = solver.step(
                                                            ^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/equinox/_module/_prebuilt.py", line 34, in __call__
    return self.__func__(self.__self__, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_solver/runge_kutta.py", line 1194, in step
    y_error = jtu.tree_map(
              ^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_solver/runge_kutta.py", line 1195, in <lambda>
    lambda _y_error: jnp.where(is_okay(result), _y_error, jnp.inf),
                               ^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_solution.py", line 54, in is_okay
    return is_successful(result) | is_event(result)
           ~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 609, in deferring_binary_op
    return binary_op(*args)
           ^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/ufunc_api.py", line 182, in __call__
    return call(*args)
           ^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/sharding.py", line 197, in shard_shape
    return _common_shard_shape(self, global_shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/sharding.py", line 64, in _common_shard_shape
    partitions, _ = get_num_ways_dim_sharded(hlo_sharding)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/op_shardings.py", line 30, in get_num_ways_dim_sharded
    assert not hlo_sharding.is_manual()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

[yashk2810]

Hi! I looked into this a bit and looks like a diffrax bug to me.

With check_vma=True, I get:

TypeError: while_loop body function carry input and carry output must have equal types, but they differ:

The input carry component x[3][2] has type float32[] but the corresponding output carry component has type float32[]{V:i}, so the varying manual axes do not match.

This might be fixed by applying `jax.lax.pcast(..., ('i',), to='varying')` to the initial carry value corresponding to the input carry component x[3][2].
See https://docs.jax.dev/en/latest/notebooks/shard_map.html#scan-vma for more information.

Revise the function so that all output types match the corresponding input types.

This should indeed be fixed by diffrax/equinox.

With check_vma=False, I get the same assert not hlo_sharding.is_manual() as you and after looking at the stack trace, there are some Manual arrays being escaped from the scope of shard_map. This shouldn't be happening and again points to a diffrax bug. If you can create a pure JAX bug, then I can take a look.