jax-ml/jax · error · ValueError

{op} does not have an {name} attribute

Error message

{op} does not have an {name} attribute

What it means

inference_utils._array_attr fetches a required MLIR array attribute (e.g. in_layouts, out_layouts, in_transforms, out_transforms, in_tmem_layouts, out_tmem_layouts) from an operation. If the op lacks the attribute entirely, a ValueError is raised naming the op and missing attribute. It typically means an operation was constructed without the layout/transform attributes the inference utilities expect.

Source

Thrown at jax/experimental/mosaic/gpu/inference_utils.py:88

    ValueError: If the operation does not have an in_tmem_layouts attribute.
  """
  return _array_attr(op, "in_tmem_layouts")


def out_tmem_layouts(op: MlirOperation) -> Sequence[ir.Attribute]:
  """Returns the out_tmem_layouts attribute of the given operation.

  Raises:
    ValueError: If the operation does not have an out_tmem_layouts attribute.
  """
  return _array_attr(op, "out_tmem_layouts")


def _array_attr(op: MlirOperation, name: str) -> Sequence[ir.Attribute]:
  try:
    result = op.attributes[name]
  except KeyError:
    raise ValueError(f"{op} does not have an {name} attribute") from None
  if not isinstance(result, ir.ArrayAttr):
    raise TypeError(f"{op} has {name} of an unexpected type: {result}")
  return result  # pyrefly: ignore[bad-return]


def should_have_in_tmem_layout(op: MlirOperation) -> bool:
  """Returns 'true' if the operation operands should be assigned a TMEM layout."""
  return any(
      isinstance(v.type, ir.MemRefType) and utils.is_tmem_ref(v)
      for v in op.operands
  )


def should_have_out_tmem_layout(op: MlirOperation) -> bool:
  """Returns 'true' if the operation results should be assigned a TMEM layout."""
  return any(
      isinstance(v.type, ir.MemRefType) and utils.is_tmem_ref(v)
      for v in op.results

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the operation is built with the required attribute (e.g. include in_layouts/out_layouts array attributes when creating the op)
2. Check op.attributes names before calling the helper and skip ops that legitimately lack them
3. Verify you are running inference on ops produced by the current Mosaic pipeline, not hand-written or stale ops from an older version

Example fix

# before
layouts = inference_utils.in_layouts(op)  # op lacks 'in_layouts' -> ValueError

# after
if "in_layouts" in op.attributes:
  layouts = inference_utils.in_layouts(op)
else:
  layouts = None  # handle op without layouts

required = ['in_layouts', 'out_layouts']
missing = [n for n in required if n not in op.attributes]
if missing:
  raise RuntimeError(f'{op.name} missing attributes: {missing}')
layouts = inference_utils.in_layouts(op)

try:
  layouts = inference_utils.in_layouts(op)
except ValueError as e:
  if 'does not have' in str(e):
    layouts = None  # op has no layouts; handle gracefully
  else:
    raise

When it happens

Trigger: Calling in_layouts(op)/out_layouts(op)/in_transforms(op)/out_transforms(op)/in_tmem_layouts(op)/out_tmem_layouts(op) on an MlirOperation missing the corresponding array attribute, e.g. a manually built mosaic op without the in_layouts attribute, or running layout inference on an op type the utilities don't cover.

Common situations: Constructing MLIR ops by hand or via text parsing without attaching all attributes; version skew where attribute names changed between Mosaic/JAX releases; calling inference helpers on ops from a different dialect.

Related errors

• {op} has {name} of an unexpected type: {result}
• Only unit strides are supported but got {op.static_strides}.
• Transposed memrefs are not supported in ExpandShapeOp.
• CollapseShapeOp with empty reassociation is not supported.
• CollapseShapeOp with non-contiguous strides is not supported

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