{"record":{"id":"9906f2c2bb54d4d6","repo":"jax-ml/jax","slug":"custom-vjp-symbolic-zeros-support-with-shard-map-i","errorCode":null,"errorMessage":"custom_vjp symbolic_zeros support with shard_map is not implemented; please open an issue at https://github.com/jax-ml/jax/issues","messagePattern":"custom_vjp symbolic_zeros support with shard_map is not implemented; please open an issue at https://github\\.com/jax-ml/jax/issues","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1463,"sourceCode":"                for mat, spec in zip(out_mats_, out_specs)]\n    return out_vals.map2(out_mats,\n                         lambda val, vma: ShardMapTracer(self, vma, val))\n\n  def process_custom_jvp_call(self, prim, fun, jvp, tracers, /, *, symbolic_zeros):\n    # Since ShardMapTrace is only used as a base main, we can drop the jvp.\n    del prim, jvp, symbolic_zeros\n    in_vals, in_mat = unzip2(map(self.to_val_mat_pair, tracers))\n    out_vals, out_mat = _run_shmap_lu(fun, self.mesh, self.manual_axes, in_vals,\n                                      in_mat, self.check)\n    return map(partial(ShardMapTracer, self), out_mat, out_vals)\n\n  def process_custom_vjp_call(self, prim, fun, fwd, bwd, tracers, /, *, out_trees,\n                              symbolic_zeros):\n    if symbolic_zeros:\n      msg = (\"custom_vjp symbolic_zeros support with shard_map is not \"\n             \"implemented; please open an issue at \"\n             \"https://github.com/jax-ml/jax/issues\")\n      raise NotImplementedError(msg)\n    del prim, fwd, bwd, out_trees, symbolic_zeros\n    in_vals, in_mat = unzip2(map(self.to_val_mat_pair, tracers))\n    out_vals, out_mat = _run_shmap_lu(fun, self.mesh, self.manual_axes, in_vals,\n                                      in_mat, self.check)\n    return map(partial(ShardMapTracer, self), out_mat, out_vals)\n\n\nclass ShardMapTracer(core.Tracer[ShardMapTrace]):\n  mat: core.ManualAxisType\n  val: JaxType\n\n  def __init__(self, trace, mat, val):\n    assert isinstance(mat, core.ManualAxisType)\n    aval = core.typeof(val)\n    mat = (mat if trace.check else\n           core.ManualAxisType(varying=trace.manual_axes))\n    size = prod(trace.mesh.shape[n] for n in mat.varying)\n    out = core.mapped_aval(size, 0, aval)","sourceCodeStart":1445,"sourceCodeEnd":1481,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1445-L1481","documentation":"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.","triggerScenarios":"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.","commonSituations":"Distributed training code using custom gradient rules with symbolic zero detection for sparsity, wrapped in shard_map for FSDP/TP style sharding.","solutions":["Set symbolic_zeros=False on the custom_vjp and handle zeros manually (check if cotangent is zero inside bwd)","Move the custom_vjp call outside the shard_map body","Open an upstream issue with the use case if the feature is required"],"exampleFix":"# before\n@custom_vjp(fun, symbolic_zeros=True)\ndef loss(...)\ny = shard_map(lambda x: loss(x), mesh, ...)(x)\n# after\n@custom_vjp(fun, symbolic_zeros=False)\ndef loss(...)\ny = shard_map(lambda x: loss(x), mesh, ...)(x)","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try: f(x) except NotImplementedError as e: if 'symbolic_zeros' in str(e): rebind custom_vjp with symbolic_zeros=False; else: raise","preventionTips":["Prefer symbolic_zeros=False in sharded training code","Check library flags for custom gradients before wrapping in shard_map"],"tags":["shard-map","custom-vjp","symbolic-zeros","autodiff","jax"],"backgroundTag":"autodiff-op-unsupported-in-sharded-region","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}