jax-ml/jax · error · NotImplementedError

Transpose of Einsum with multiple linear inputs is not suppo

Error message

Transpose of Einsum with multiple linear inputs is not supported.

What it means

Raised by Einsum.transpose (the VJP rule) when gradient accumulation is requested for more than one input simultaneously. The linearized transpose implementation only supports a single GradAccum per einsum; multiple linear inputs under custom accumulation (e.g. jax.checkpoint/accumulation APIs) are unimplemented.

Source

Thrown at jax/_src/numpy/hijax.py:443

    return batched_prim(*args), 0

  def jvp(self, primals: tuple[Array, ...], tangents: Any) -> tuple[Array, Array]:
    primal_out = self(*primals)
    tangent_outs = []
    for i, t in enumerate(tangents):
      if not isinstance(t, ad_util.Zero):
        tangent_outs.append(self(*primals[:i], t, *primals[i+1:]))
    if not tangent_outs:
      return primal_out, ad_util.zeros_like_aval(self.out_aval)
    return primal_out, functools.reduce(lax.add, tangent_outs)

  def transpose(self, out_ct, *maybe_accums):
    in_subs_list = self.subscripts.split('->')[0].split(',')
    out_sub = self.subscripts.split('->')[1]

    accums = [acc for acc in maybe_accums if isinstance(acc, ad.GradAccum)]
    if len(accums) > 1:
      raise NotImplementedError("Transpose of Einsum with multiple linear inputs is not supported.")

    for i, accum in enumerate(maybe_accums):
      if isinstance(accum, ad.GradAccum):
        if isinstance(out_ct, ad_util.Zero):
          accum.accum(ad_util.zeros_like_aval(self.in_avals[i]))
          continue

        orig_sub = in_subs_list[i]
        orig_aval = self.in_avals[i]

        all_ct_input_chars = set(out_sub)
        for k, sub in enumerate(in_subs_list):
          if k != i:
            all_ct_input_chars.update(sub)

        missing_chars = sorted(
            set(orig_sub) - all_ct_input_chars, key=orig_sub.index
        )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Accumulate gradients for only one einsum input at a time (run separate VJPs per input)
2. Avoid gradient accumulation around this einsum: remove jax.checkpoint/custom accum on that call or split the einsum so each has one linear input
3. File/track an upstream issue in JAX for multi-input accumulation support

Example fix

# before
pull = jax.vjp(lambda a, b: jnp.einsum('ij,jk->ik', a, b), a, b)
# with accumulation on both inputs -> NotImplementedError
# after: accumulate one input at a time
_, vjp_a = jax.vjp(lambda a: jnp.einsum('ij,jk->ik', a, b_stop), a)
_, vjp_b = jax.vjp(lambda b: jnp.einsum('ij,jk->ik', a_stop, b), b)

try:
    grads = jax.grad(loss)(params)
except NotImplementedError as e:
    if 'multiple linear inputs' in str(e):
        # differentiate one einsum input at a time
        grads = {k: jax.grad(lambda p: loss_with_fixed(others, k, p))(p) for k, p in params.items()}
    else:
        raise

When it happens

Trigger: Computing a VJP of an einsum with two or more linear operands while using gradient accumulation (ad.GradAccum), e.g. inside remat/checkpointed code with per-input accumulation; the trace reaches _nonzero_impl-adjacent transpose machinery and hits this NotImplementedError.

Common situations: Backprop through large einsum-based models under jax.checkpoint or custom gradient accumulation; upgrading code that previously differentiated without accumulation.

Related errors

• 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
• for transpose support, subclass {type(self)} must implement
• open an issue at https://github.com/google/jax !!

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