jax-ml/jax · error · NotImplementedError

for transpose support, subclass {type(self)} must implement

Error message

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

What it means

HiPrim's optional transpose rule is a stub. Primitives that are linear in some inputs can implement `transpose(out_ct, *maybe_accums)` to support transformations like linear_transpose; the default raises.

Source

Thrown at jax/_src/hijax.py:207

  def jvp(self, primals, tangents):
    raise NotImplementedError(f"for jvp support, subclass {type(self)} must "
                              "implement `jvp`")

  def lin(self, nzs_in, *primals):
    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`")

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Implement `def transpose(self, out_ct, *maybe_accums)` returning input cotangents (or None plus a log dict)
2. Prefer vjp_from_jvp / vjp_from_lin derivations which supply transposes automatically
3. Avoid linear_transpose over this primitive

Example fix

class LinearPrim(hijax.HiPrim):
  # after
  def transpose(self, out_ct, *args):
    return tree_map(lambda x: x, args)  # rule matching the linear map

if type(prim).transpose is hijax.HiPrim.transpose:
    raise ValueError(f'{type(prim).__name__} lacks a transpose rule')

def has_transpose_rule(p) -> bool:
    return type(p).transpose is not hijax.HiPrim.transpose

try:
    jax.linear_transpose(f, x)(y)
except NotImplementedError as e:
    if 'transpose' in str(e):
        raise RuntimeError('use vjp instead of linear_transpose for this primitive') from e
    raise

When it happens

Trigger: jax.linear_transpose (or fancy transpose machinery, _call_hi_primitive_transpose) reaching a HiPrim subclass that did not override transpose.

Common situations: Using jax.linear_transpose or a VJP derivation path that transposes the primitive when only nonlinear rules were defined.

Related errors

• Pure callbacks do not support transpose. Please use `jax.cus
• for grad support, subclass {type(self)} must implement `vjp_
• for jvp support, subclass {type(self)} must implement `jvp`
• for linearize support, subclass {type(self)} must implement
• {type(_prim).__name__}.transpose should return None or a dic

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