{"record":{"id":"e4f662de4aafa016","repo":"jax-ml/jax","slug":"transpose-with-implicit-broadcasting-of-unshaped-v","errorCode":null,"errorMessage":"transpose with implicit broadcasting of unshaped values. Got {type(aval)}","messagePattern":"transpose with implicit broadcasting of unshaped values\\. Got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":4441,"sourceCode":"  dtype_rule = partial(naryop_dtype_rule, result_dtype, accepted_dtypes, name,\n                       allow_extended_dtype=allow_extended_dtype,\n                       require_same=require_same_dtypes)\n  shape_rule = partial(broadcasting_shape_rule, name)\n  sharding_rule = partial(broadcasting_sharding_rule, name)\n  prim = standard_primitive(\n      shape_rule, dtype_rule, name, sharding_rule=sharding_rule,\n      vma_rule=partial(core.standard_vma_rule, name),\n      ur_rule=partial(nary_ur_rule, name) if ur_rule is None else ur_rule)\n  batching.defbroadcasting(prim)\n  return prim\nstandard_naryop = partial(naryop, input_dtype)\n\n\n# Like autograd.numpy.numpy_vjps.unbroadcast, this utility handles transposition\n# involving linear primitives with implicit broadcasting.\ndef _unbroadcast(aval, x):\n  if not isinstance(aval, ShapedArray):\n    raise TypeError(\n        'transpose with implicit broadcasting of unshaped values. Got'\n        f' {type(aval)}')\n  x_shape = np.shape(x)\n  if (core.definitely_equal_shape(aval.shape, x_shape) and\n      aval.sharding == typeof(x).sharding):\n    return x\n  assert not aval.shape or len(x_shape) == len(aval.shape)\n  if not aval.shape:\n    return reduce_sum(x, list(range(len(x_shape))))\n  else:\n    dims = [i for i, (a, b) in enumerate(zip(x_shape, aval.shape))\n            if not core.definitely_equal(a, b)]\n    if config.enable_checks.value:\n      assert all(aval.shape[i] == 1 for i in dims)\n    x = reduce_sum(x, dims) if dims else x\n    return reshape(x, aval.shape, out_sharding=aval.sharding)\n\ndef _maybe_broadcast(target_shape, x, target_sharding):","sourceCodeStart":4423,"sourceCodeEnd":4459,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L4423-L4459","documentation":"_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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Simplify to static shapes for the differentiated region (avoid dynamic shape polymorphism around grad)","Rewrite so broadcasting happens outside the differentiated linear op (pre-broadcast inputs)","Update JAX — support for unbroadcasting dynamic-shaped avals improves across versions; if it persists, file an issue"],"exampleFix":"// before\n# dynamic-shape input + grad through implicit broadcast\nf = jax.grad(lambda x: lax.broadcast_in_dim(x, dyn_shape, (0,)).sum())\n// after\nf = jax.grad(lambda x: jnp.broadcast_to(x[:, None], fixed_shape).sum())","handlingStrategy":"fallback","validationCode":"null","typeGuard":null,"tryCatchPattern":"try:\n    g = jax.grad(f)(x)\nexcept TypeError as e:\n    if 'unshaped values' in str(e):\n        g = jax.grad(f_static)(x)  # static-shape reimplementation\n    else:\n        raise","preventionTips":["Avoid dynamic/shape-polymorphic inputs inside grad on broadcasting linear ops","Pre-broadcast operands explicitly instead of relying on implicit broadcasting under vjp"],"tags":["jax","autodiff","vjp","dynamic-shapes","transpose","type-error"],"backgroundTag":"autodiff-unsupported-aval","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}