jax-ml/jax · error · TypeError

{type(_prim).__name__}.transpose should return None or a dic

Error message

{type(_prim).__name__}.transpose should return None or a dict of backward-pass log entries, got {type(log).__name__}

What it means

The fancy-transpose handler for call_hi_primitive_p calls _prim.transpose(cts, *primals); its return must be None or a dict of log entries. The actual transposed values are delivered through the machinery, so returning cotangents directly is a TypeError.

Source

Thrown at jax/_src/hijax.py:520

    del params['has_sres']
  return core._pp_eqn(eqn.replace(params=params), context, settings)
core.pp_eqn_rules[call_hi_primitive_linearized_p] = _call_hi_primitive_linearized_prettyprint

def _call_hi_primitive_jvp(primals, tangents, *, _prim):
  primals = tree_unflatten(_prim.in_tree, primals)
  tangents = tree_unflatten(_prim.in_tree, tangents)
  out_primals, out_tangents = _prim.jvp(primals, tangents)
  out_primals_flat = tree_leaves_checked(_prim.out_tree, out_primals)
  out_tangents_flat = _prim.out_tree.flatten_up_to(out_tangents)
  return out_primals_flat, out_tangents_flat
ad.primitive_jvps[call_hi_primitive_p] = _call_hi_primitive_jvp

def _call_hi_primitive_transpose(cts_flat, *primals_flat, _prim):
  cts = tree_unflatten(_prim.out_tree, cts_flat)
  primals = tree_unflatten(_prim.in_tree, primals_flat)
  log = _prim.transpose(cts, *primals)  # a returned dict logs entries
  if log is not None and type(log) is not dict:
    raise TypeError(
        f"{type(_prim).__name__}.transpose should return None or a dict of "
        f"backward-pass log entries, got {type(log).__name__}")
  return log
ad.fancy_transposes[call_hi_primitive_p] = _call_hi_primitive_transpose

def _call_hi_primitive_dce(used_outs_flat, eqn):
  _prim = eqn.params['_prim']
  used_out = tree_unflatten(_prim.out_tree, used_outs_flat)
  used_ins, produced_outs, new_prim = _prim.dce(used_out)
  if new_prim is None:
    return [False] * len(eqn.invars), None
  name = f'{type(_prim).__name__}.dce'
  used_ins_flat = api.tuptree_flags(
      used_ins, _prim.in_tree, 'used_ins',
      f'the first (used inputs) return value of {name}')
  produced_outs_flat = api.tuptree_flags(
      produced_outs, _prim.out_tree, 'produced_outs',
      f'the second (produced outputs) return value of {name}')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rework transpose to pass results via its output-accumulator mechanism and return None or a logs dict
2. Follow the HiPrim.transpose docstring/examples for the expected signature `transpose(self, out_ct, *maybe_accums)`

Example fix

# before
def transpose(self, out_ct, *args):
    return args
# after
def transpose(self, out_ct, *args):
    ...
    return None

def transpose_returns_valid(p) -> bool:
    out = p.transpose(dummy_ct, *dummy_primals)
    return out is None or type(out) is dict

try:
    jax.linear_transpose(f, x)(y)
except TypeError as e:
    if 'transpose should return' in str(e):
        raise RuntimeError('transpose must return None or a logs dict') from e
    raise

When it happens

Trigger: A custom HiPrim transpose implementation returns cotangents (tuple/array) instead of None or a logs dict.

Common situations: Writing transpose with the conventional JAX convention (return input cotangents) rather than hijax's logging convention.

Related errors

• Pure callbacks do not support transpose. Please use `jax.cus
• for transpose support, subclass {type(self)} must implement
• {type(_prim).__name__} returned structured residuals from `v
• {type(_prim).__name__}.vjp_bwd should return None or a dict
• transpose_solve required for backwards mode automatic differ

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