jax-ml/jax · error · NotImplementedError

for vmap support, subclass {type(self)} must implement `batc

Error message

for vmap support, subclass {type(self)} must implement `batch` or `batch_dim_rule`

What it means

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.

Source

Thrown at jax/_src/hijax.py:216

  def linearized(self, residuals, *tangents):
    raise NotImplementedError(
        f"for linearize support, subclass {type(self)} must implement `lin` "
        "and `linearized`, or derive them from its `jvp` rule by setting "
        "`lin, linearized = linearize_from_jvp`")

  # optional transpose rule, for primitives that are linear in some inputs
  def transpose(self, out_ct, *maybe_accums):
    raise NotImplementedError(f"for transpose support, subclass {type(self)} "
                              "must implement `transpose`")

  # vmap interface
  def batch(self, axis_data, args, dims):
    out_dim = self.batch_dim_rule(axis_data, dims)
    return VmapOf(self, axis_data, dims, out_dim)(*args), out_dim

  def batch_dim_rule(self, axis_data, dims, /):
    raise NotImplementedError(f"for vmap support, subclass {type(self)} must "
                              "implement `batch` or `batch_dim_rule`")

  # optional dce control
  def dce(self, used_outs):
    used_outs_flat = tree_leaves_checked(self.out_tree, used_outs)
    if not any(used_outs_flat):
      return False, False, None
    else:
      return True, True, self

  # optional remat control
  def remat(self, _trace, *args):
    return self(*args), self  # full remat by default

  def __call__(self, *args):
    args_flat = tree_leaves_checked(self.in_tree, args)
    ans_flat = call_hi_primitive_p.bind(*args_flat, _prim=self)
    return tree_unflatten(self.out_tree, ans_flat)

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

class MyPrim(hijax.HiPrim):
  # after
  def batch_dim_rule(self, axis_data, dims):
    return tuple(d + 1 if d is not None else None for d in dims)

Defensive patterns

Strategy: validation

Validate before calling

if (type(prim).batch is hijax.HiPrim.batch and
        type(prim).batch_dim_rule is hijax.HiPrim.batch_dim_rule):
    raise ValueError(f'{type(prim).__name__} lacks vmap rules')

Type guard

def has_vmap_rules(p) -> bool:
    return not (type(p).batch is hijax.HiPrim.batch and
                type(p).batch_dim_rule is hijax.HiPrim.batch_dim_rule)

Try / catch

try:
    jax.vmap(f)(xs)
except NotImplementedError as e:
    if 'vmap' in str(e):
        return jax.lax.map(f, xs)  # sequential fallback
    raise

Prevention

Implement batch_dim_rule in every HiPrim subclass
Test each primitive under jax.vmap with batched and unbatched dims

When it happens

Trigger: 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.

Common situations: Custom primitive works scalar-wise, then the model is batched for training/inference under vmap.

Related errors

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