{"record":{"id":"2f03168293d000ed","repo":"jax-ml/jax","slug":"for-vmap-support-subclass-type-self-must-imple","errorCode":null,"errorMessage":"for vmap support, subclass {type(self)} must implement `batch` or `batch_dim_rule`","messagePattern":"for vmap support, subclass (.+?) must implement `batch` or `batch_dim_rule`","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":216,"sourceCode":"\n  def linearized(self, residuals, *tangents):\n    raise NotImplementedError(\n        f\"for linearize support, subclass {type(self)} must implement `lin` \"\n        \"and `linearized`, or derive them from its `jvp` rule by setting \"\n        \"`lin, linearized = linearize_from_jvp`\")\n\n  # optional transpose rule, for primitives that are linear in some inputs\n  def transpose(self, out_ct, *maybe_accums):\n    raise NotImplementedError(f\"for transpose support, subclass {type(self)} \"\n                              \"must implement `transpose`\")\n\n  # vmap interface\n  def batch(self, axis_data, args, dims):\n    out_dim = self.batch_dim_rule(axis_data, dims)\n    return VmapOf(self, axis_data, dims, out_dim)(*args), out_dim\n\n  def batch_dim_rule(self, axis_data, dims, /):\n    raise NotImplementedError(f\"for vmap support, subclass {type(self)} must \"\n                              \"implement `batch` or `batch_dim_rule`\")\n\n  # optional dce control\n  def dce(self, used_outs):\n    used_outs_flat = tree_leaves_checked(self.out_tree, used_outs)\n    if not any(used_outs_flat):\n      return False, False, None\n    else:\n      return True, True, self\n\n  # optional remat control\n  def remat(self, _trace, *args):\n    return self(*args), self  # full remat by default\n\n  def __call__(self, *args):\n    args_flat = tree_leaves_checked(self.in_tree, args)\n    ans_flat = call_hi_primitive_p.bind(*args_flat, _prim=self)\n    return tree_unflatten(self.out_tree, ans_flat)","sourceCodeStart":198,"sourceCodeEnd":234,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L198-L234","documentation":"HiPrim's vmap support requires either a `batch` method or a `batch_dim_rule(axis_data, dims)`; the default batch() calls batch_dim_rule, which raises if neither is overridden.","triggerScenarios":"Calling jax.vmap (or a transform that internally vmaps, like pmap or batched solvers) on a function applying a HiPrim subclass with no batching rules.","commonSituations":"Custom primitive works scalar-wise, then the model is batched for training/inference under vmap.","solutions":["Implement `def batch_dim_rule(self, axis_data, dims)` returning output dims (default batch() then re-applies the primitive vmapped)","Or fully override `def batch(self, axis_data, args, dims)`","Test the primitive under jax.vmap as part of its unit tests"],"exampleFix":"class MyPrim(hijax.HiPrim):\n  # after\n  def batch_dim_rule(self, axis_data, dims):\n    return tuple(d + 1 if d is not None else None for d in dims)\n","handlingStrategy":"validation","validationCode":"if (type(prim).batch is hijax.HiPrim.batch and\n        type(prim).batch_dim_rule is hijax.HiPrim.batch_dim_rule):\n    raise ValueError(f'{type(prim).__name__} lacks vmap rules')","typeGuard":"def has_vmap_rules(p) -> bool:\n    return not (type(p).batch is hijax.HiPrim.batch and\n                type(p).batch_dim_rule is hijax.HiPrim.batch_dim_rule)","tryCatchPattern":"try:\n    jax.vmap(f)(xs)\nexcept NotImplementedError as e:\n    if 'vmap' in str(e):\n        return jax.lax.map(f, xs)  # sequential fallback\n    raise","preventionTips":["Implement batch_dim_rule in every HiPrim subclass","Test each primitive under jax.vmap with batched and unbatched dims"],"tags":["jax","vmap","batching","custom-primitive"],"backgroundTag":"missing-vmap-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}