{"record":{"id":"2803b5bae2ae21c7","repo":"jax-ml/jax","slug":"cannot-splat-a-tensor","errorCode":null,"errorMessage":"cannot splat a tensor","messagePattern":"cannot splat a tensor","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1543,"sourceCode":"def _zeros(t: ir.Type) -> ir.Value:\n  return _full(t, 0)\n\n\ndef _zeros_like(x: ir.Value) -> ir.Value:\n  return _full(x.type, 0)\n\n\ndef _ones(t: ir.Type) -> ir.Value:\n  return _full(t, 1)\n\n\ndef _ones_like(x: ir.Value) -> ir.Value:\n  return _full(x.type, 1)\n\n\ndef _splat(x: ir.Value, shape: Sequence[int]) -> ir.Value:\n  if isinstance(x.type, ir.RankedTensorType):\n    raise TypeError(\"cannot splat a tensor\")\n  if not shape:\n    return x\n  return tt_dialect.splat(ir.RankedTensorType.get(shape, x.type), x)\n\n\ndef _expand_dims(x: ir.Value, axis: int) -> ir.Value:\n  if not isinstance(x.type, ir.RankedTensorType):\n    shape = list(ir.RankedTensorType(x.type).shape)\n    shape.insert(axis, 1)\n    return _splat(x, shape)\n  return tt_dialect.expand_dims(x, axis)\n\n\ndef _float_float_cast(src: ir.Value, dst_type: ir.Type) -> ir.Value:\n  src_element_type = ir.FloatType(_element_type(src.type))\n  dst_element_type = ir.FloatType(_element_type(dst_type))\n  if src_element_type.width == 8 or dst_element_type.width == 8:\n    rounding = (","sourceCodeStart":1525,"sourceCodeEnd":1561,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1525-L1561","documentation":"_splat wraps Triton's splat op, which broadcasts a scalar into a tensor of the given shape; it type-errors if handed a value that is already a RankedTensorType. It is reached from _expand_dims and _reshape when those helpers need to broadcast a value that should be scalar but is actually a tensor.","triggerScenarios":"Internal lowering paths (_expand_dims/_reshape) receiving a tensor-typed ir.Value where a scalar is expected — typically a consequence of earlier lowering producing a 0-d or 1-d tensor instead of a scalar, e.g. when reshaping a block to a scalar shape or expanding dims of an already-tensorized index value.","commonSituations":"Pallas kernels that use scalar-shaped blocks (shape ()) or reshape blocks down to scalars; grid/block index arithmetic that yields tensor-typed values; version mismatches between jax and jax-triton where scalar handling changed.","solutions":["Reshape the kernel logic to keep values tensor-shaped (minimum block dim of 1) rather than scalar-shaped; avoid jnp.reshape(x, ()) inside Triton kernels","Upgrade jax (and jax-triton if installed separately) to a version where scalar block handling in expand_dims/reshape lowering is fixed","If you control the input, convert the tensor to a scalar before the operation (e.g. extract the element) so _splat sees a scalar"],"exampleFix":"// before\ns = kernel_ref.reshape(())  # scalar-shaped block hits _splat with tensor\n// after\ns = kernel_ref.reshape((1,))  # keep 1-element tensor shape","handlingStrategy":"validation","validationCode":"def check_block_shapes(shapes):\n    assert all(len(s) > 0 and all(d >= 1 for d in s) for s in shapes), \\\n        'avoid scalar-shaped (() ) blocks; use (1,) instead'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never reshape blocks to shape (); keep at least one dimension of size 1","Pin compatible jax/jax-triton versions where scalar reshape lowering is fixed"],"tags":["jax","pallas","triton","splat","scalar-shape","reshape","typeerror"],"backgroundTag":"invalid-shape-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}