{"record":{"id":"76f07b2854b26f1a","repo":"jax-ml/jax","slug":"for-transpose-support-subclass-type-self-must","errorCode":null,"errorMessage":"for transpose support, subclass {type(self)} must implement `transpose`","messagePattern":"for transpose support, subclass (.+?) must implement `transpose`","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":207,"sourceCode":"  def jvp(self, primals, tangents):\n    raise NotImplementedError(f\"for jvp support, subclass {type(self)} must \"\n                              \"implement `jvp`\")\n\n  def lin(self, nzs_in, *primals):\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  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","sourceCodeStart":189,"sourceCodeEnd":225,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L189-L225","documentation":"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.","triggerScenarios":"jax.linear_transpose (or fancy transpose machinery, _call_hi_primitive_transpose) reaching a HiPrim subclass that did not override transpose.","commonSituations":"Using jax.linear_transpose or a VJP derivation path that transposes the primitive when only nonlinear rules were defined.","solutions":["Implement `def transpose(self, out_ct, *maybe_accums)` returning input cotangents (or None plus a log dict)","Prefer vjp_from_jvp / vjp_from_lin derivations which supply transposes automatically","Avoid linear_transpose over this primitive"],"exampleFix":"class LinearPrim(hijax.HiPrim):\n  # after\n  def transpose(self, out_ct, *args):\n    return tree_map(lambda x: x, args)  # rule matching the linear map","handlingStrategy":"validation","validationCode":"if type(prim).transpose is hijax.HiPrim.transpose:\n    raise ValueError(f'{type(prim).__name__} lacks a transpose rule')","typeGuard":"def has_transpose_rule(p) -> bool:\n    return type(p).transpose is not hijax.HiPrim.transpose","tryCatchPattern":"try:\n    jax.linear_transpose(f, x)(y)\nexcept NotImplementedError as e:\n    if 'transpose' in str(e):\n        raise RuntimeError('use vjp instead of linear_transpose for this primitive') from e\n    raise","preventionTips":["Only promise transpose for genuinely linear primitives","Prefer vjp_from_jvp which supplies transposes"],"tags":["jax","autodiff","transpose","not-implemented"],"backgroundTag":"missing-autodiff-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}