jax-ml/jax · error · TypeError

transpose with implicit broadcasting of unshaped values. Got

Error message

transpose with implicit broadcasting of unshaped values. Got {type(aval)}

What it means

_unbroadcast is used by transpose (VJP) rules of linear primitives to sum broadcast dimensions out of a cotangent. It requires the primal's aval to be a ShapedArray; unshaped/abstract values (e.g. DShapedArray or other avals) cannot be unbroadcast, raising TypeError.

Source

Thrown at jax/_src/lax/lax.py:4441

  dtype_rule = partial(naryop_dtype_rule, result_dtype, accepted_dtypes, name,
                       allow_extended_dtype=allow_extended_dtype,
                       require_same=require_same_dtypes)
  shape_rule = partial(broadcasting_shape_rule, name)
  sharding_rule = partial(broadcasting_sharding_rule, name)
  prim = standard_primitive(
      shape_rule, dtype_rule, name, sharding_rule=sharding_rule,
      vma_rule=partial(core.standard_vma_rule, name),
      ur_rule=partial(nary_ur_rule, name) if ur_rule is None else ur_rule)
  batching.defbroadcasting(prim)
  return prim
standard_naryop = partial(naryop, input_dtype)


# Like autograd.numpy.numpy_vjps.unbroadcast, this utility handles transposition
# involving linear primitives with implicit broadcasting.
def _unbroadcast(aval, x):
  if not isinstance(aval, ShapedArray):
    raise TypeError(
        'transpose with implicit broadcasting of unshaped values. Got'
        f' {type(aval)}')
  x_shape = np.shape(x)
  if (core.definitely_equal_shape(aval.shape, x_shape) and
      aval.sharding == typeof(x).sharding):
    return x
  assert not aval.shape or len(x_shape) == len(aval.shape)
  if not aval.shape:
    return reduce_sum(x, list(range(len(x_shape))))
  else:
    dims = [i for i, (a, b) in enumerate(zip(x_shape, aval.shape))
            if not core.definitely_equal(a, b)]
    if config.enable_checks.value:
      assert all(aval.shape[i] == 1 for i in dims)
    x = reduce_sum(x, dims) if dims else x
    return reshape(x, aval.shape, out_sharding=aval.sharding)

def _maybe_broadcast(target_shape, x, target_sharding):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Simplify to static shapes for the differentiated region (avoid dynamic shape polymorphism around grad)
2. Rewrite so broadcasting happens outside the differentiated linear op (pre-broadcast inputs)
3. Update JAX — support for unbroadcasting dynamic-shaped avals improves across versions; if it persists, file an issue

Example fix

// before
# dynamic-shape input + grad through implicit broadcast
f = jax.grad(lambda x: lax.broadcast_in_dim(x, dyn_shape, (0,)).sum())
// after
f = jax.grad(lambda x: jnp.broadcast_to(x[:, None], fixed_shape).sum())

null

try:
    g = jax.grad(f)(x)
except TypeError as e:
    if 'unshaped values' in str(e):
        g = jax.grad(f_static)(x)  # static-shape reimplementation
    else:
        raise

When it happens

Trigger: Reverse-mode differentiation through a linear lax op whose primal aval is not a ShapedArray — typically dynamic-shape arrays (DShapedArray) inside grad/vjp with implicit broadcasting.

Common situations: Using dynamic shapes (jnp.eager tracing / shape polymorphism) with grad on broadcasting linear ops like broadcast_in_dim, reshape, or pad; newer JAX versions where dynamic-shaped avals flow into VJP rules that predate them.

Related errors

• Pure callbacks do not support transpose. Please use `jax.cus
• primal and tangent arguments to jax.jvp must be tuples or li
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• structure of the differentiated function {jaxpr.debug_info.f

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