{"record":{"id":"b2973dfa876124f7","repo":"jax-ml/jax","slug":"only-arguments-with-shape-1-are-supported","errorCode":null,"errorMessage":"Only arguments with shape [..., 1] are supported.","messagePattern":"Only arguments with shape \\[\\.\\.\\., 1\\] are supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":1826,"sourceCode":"\ndef get_join_type(old_type: ir.RankedTensorType):\n  shape = old_type.shape\n  shape.append(2)\n  return ir.RankedTensorType.get(shape, old_type.element_type, old_type.encoding)\n\n\n@register_lowering(lax.concatenate_p)\ndef _concatenate_lowering_rule(ctx: LoweringRuleContext, *args, dimension):\n  if len(args) != 2:\n    raise NotImplementedError(\"Only 2-argument concatenate is supported.\")\n  x_aval, y_aval = ctx.avals_in\n  x, y = args\n  if dimension != x_aval.ndim-1:\n    raise NotImplementedError(\n        \"Only concatenate along the last dimension is supported.\"\n    )\n  if x_aval.shape[-1] != 1 or y_aval.shape[-1] != 1:\n    raise NotImplementedError(\n        \"Only arguments with shape [..., 1] are supported.\"\n    )\n  lhs = _reshape(x, x_aval.shape[:-1])\n  rhs = _reshape(y, y_aval.shape[:-1])\n  ret_type = get_join_type(ir.RankedTensorType(rhs.type))\n  return tt_dialect.join(ret_type, lhs, rhs)\n\n@register_lowering(jax._src.lax.lax.stack_p)\ndef _stack_lowering_rule(ctx: LoweringRuleContext, *args, axis):\n  if len(args) != 2:\n    raise NotImplementedError(\"Only 2-argument stack is supported in Triton.\")\n  [x_aval, y_aval] = ctx.avals_in\n  x, y = args\n  if axis != x_aval.ndim:\n    raise NotImplementedError(\"Only stack along the last dimension is supported in Triton.\")\n\n  x = _ensure_ir_value(x, x_aval)\n  y = _ensure_ir_value(y, y_aval)","sourceCodeStart":1808,"sourceCodeEnd":1844,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L1808-L1844","documentation":"The lowering of concatenate to tt.join requires each argument to have shape [..., 1] — i.e. the last dimension must be exactly 1 so it can be reshaped away and re-joined as a new minor dimension of size 2. Any other trailing shape fails.","triggerScenarios":"jnp.concatenate([x, y], axis=-1) inside a Triton Pallas kernel where x.shape[-1] or y.shape[-1] is not 1, e.g. joining two [B, N] blocks with N > 1.","commonSituations":"Trying to append a column of values to a block (shape [B, k] + [B, k]) rather than pairing scalars; generic concatenation assumed from numpy semantics.","solutions":["Restructure to pair singleton trailing dims: expand_dims each operand to [..., 1] before concatenating","Do general concatenation outside the kernel on the host/XLA side","Build the combined block with explicit indexing/store into a preallocated reference instead of concatenate"],"exampleFix":"// before\nout = jnp.concatenate([x, y], axis=-1)  # x, y have shape [B, k]\n\n// after\nout = jnp.concatenate([x[..., None, :], y[..., None, :]], axis=-2)  # or concatenate outside kernel","handlingStrategy":"validation","validationCode":"assert x.shape[-1] == 1 and y.shape[-1] == 1, 'operands must be [..., 1] for in-kernel concat'","typeGuard":"def singleton_trailing(a) -> bool:\n    return a.shape[-1] == 1","tryCatchPattern":"try:\n    out = jnp.concatenate([x, y], axis=-1)\nexcept NotImplementedError:\n    out = jnp.stack([x, y], axis=-1)  # if pairing blocks","preventionTips":["Design kernels to pair scalars ([...,1]) rather than append columns","Use expand_dims before in-kernel concatenation","Do wide concatenation outside the kernel"],"tags":["jax","triton","pallas","concatenate","shape","not-implemented"],"backgroundTag":"unsupported-operation-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}