jax-ml/jax · error · NotImplementedError

custom_vjp symbolic_zeros support with shard_map is not impl

Error message

custom_vjp symbolic_zeros support with shard_map is not implemented; please open an issue at https://github.com/jax-ml/jax/issues

What it means

Inside a shard_map-traced region, custom_vjp calls that request symbolic_zeros cannot be handled: the interplay of zero-preserving custom gradients with manual sharding tracing is unimplemented and raises NotImplementedError.

Source

Thrown at jax/_src/shard_map.py:1463

                for mat, spec in zip(out_mats_, out_specs)]
    return out_vals.map2(out_mats,
                         lambda val, vma: ShardMapTracer(self, vma, val))

  def process_custom_jvp_call(self, prim, fun, jvp, tracers, /, *, symbolic_zeros):
    # Since ShardMapTrace is only used as a base main, we can drop the jvp.
    del prim, jvp, symbolic_zeros
    in_vals, in_mat = unzip2(map(self.to_val_mat_pair, tracers))
    out_vals, out_mat = _run_shmap_lu(fun, self.mesh, self.manual_axes, in_vals,
                                      in_mat, self.check)
    return map(partial(ShardMapTracer, self), out_mat, out_vals)

  def process_custom_vjp_call(self, prim, fun, fwd, bwd, tracers, /, *, out_trees,
                              symbolic_zeros):
    if symbolic_zeros:
      msg = ("custom_vjp symbolic_zeros support with shard_map is not "
             "implemented; please open an issue at "
             "https://github.com/jax-ml/jax/issues")
      raise NotImplementedError(msg)
    del prim, fwd, bwd, out_trees, symbolic_zeros
    in_vals, in_mat = unzip2(map(self.to_val_mat_pair, tracers))
    out_vals, out_mat = _run_shmap_lu(fun, self.mesh, self.manual_axes, in_vals,
                                      in_mat, self.check)
    return map(partial(ShardMapTracer, self), out_mat, out_vals)


class ShardMapTracer(core.Tracer[ShardMapTrace]):
  mat: core.ManualAxisType
  val: JaxType

  def __init__(self, trace, mat, val):
    assert isinstance(mat, core.ManualAxisType)
    aval = core.typeof(val)
    mat = (mat if trace.check else
           core.ManualAxisType(varying=trace.manual_axes))
    size = prod(trace.mesh.shape[n] for n in mat.varying)
    out = core.mapped_aval(size, 0, aval)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set symbolic_zeros=False on the custom_vjp and handle zeros manually (check if cotangent is zero inside bwd)
2. Move the custom_vjp call outside the shard_map body
3. Open an upstream issue with the use case if the feature is required

Example fix

# before
@custom_vjp(fun, symbolic_zeros=True)
def loss(...)
y = shard_map(lambda x: loss(x), mesh, ...)(x)
# after
@custom_vjp(fun, symbolic_zeros=False)
def loss(...)
y = shard_map(lambda x: loss(x), mesh, ...)(x)

try: f(x) except NotImplementedError as e: if 'symbolic_zeros' in str(e): rebind custom_vjp with symbolic_zeros=False; else: raise

When it happens

Trigger: A function decorated with custom_vjp(..., symbolic_zeros=True) (or a library that sets it, e.g. some optimizers/losses) called within a shard_map body.

Common situations: Distributed training code using custom gradient rules with symbolic zero detection for sparsity, wrapped in shard_map for FSDP/TP style sharding.

Related errors

• Pure callbacks do not support transpose. Please use `jax.cus
• at {keystr(path)}, got fwd output type {ty.str_short()} whic
• {result}the bwd rule attached to {primal_sourceinfo} produce
• No VJP defined for custom_vjp function {self.f.__name__} usi
• primal and tangent arguments to jax.jvp must be tuples or li

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/9906f2c2bb54d4d6. Report an issue: GitHub.