jax-ml/jax · error · ValueError

Unknown attribute: {attr_name}

Error message

Unknown attribute: {attr_name}

What it means

_in_attr_for_operand dispatches on the attribute name to find which operand index corresponds to a layout/transform attribute. It only knows in_layouts, out_layouts, in_transforms, out_transforms, and in_tmem_layouts; any other attr_name raises ValueError('Unknown attribute'). It indicates a programming error in the caller (misspelled or unsupported attribute name).

Source

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

  return attr[index]


def _in_attr_for_operand(
    op: MlirOperation,
    operand: ir.Value,
    attr_name: str,
) -> ir.Attribute | None:
  if attr_name == "in_layouts":
    predicate = lambda v: isinstance(v.type, ir.VectorType)
  elif attr_name == "in_transforms":
    predicate = is_transformable_smem_memref
  elif attr_name == "in_tmem_layouts":
    predicate = (
        lambda v: isinstance(v.type, ir.MemRefType)
        and ir.MemRefType(v.type).memory_space == utils.tmem()
    )
  else:
    raise ValueError(f"Unknown attribute: {attr_name}")

  operand_number = [o for o in op.operands if predicate(o)].index(operand)

  return attr_element(attr_name, op, operand_number)


in_layout_for_operand = partial(
    _in_attr_for_operand, attr_name="in_layouts"
)
in_tmem_layout_for_operand = partial(
    _in_attr_for_operand, attr_name="in_tmem_layouts"
)
in_transforms_for_operand = partial(
    _in_attr_for_operand, attr_name="in_transforms"
)


def should_have_in_transforms(op: ir.OpView) -> bool:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the dispatch branches in inference_utils.py and only pass supported attribute names (in_layouts, out_layouts, in_transforms, out_transforms, in_tmem_layouts)
2. Fix typos in the attribute name string at the call site
3. Upgrade JAX/Mosaic if you need an attribute (e.g. out_tmem_layouts) supported in a newer release

Example fix

# before
attr = inference_utils._in_attr_for_operand('in_layout', op, operand)

# after
attr = inference_utils._in_attr_for_operand('in_layouts', op, operand)

SUPPORTED = {'in_layouts', 'out_layouts', 'in_transforms', 'out_transforms', 'in_tmem_layouts'}
assert attr_name in SUPPORTED, f'unsupported attr_name: {attr_name}'
result = inference_utils._in_attr_for_operand(attr_name, op, operand)

def is_supported_attr(name: str) -> bool:
  return name in {'in_layouts', 'out_layouts', 'in_transforms', 'out_transforms', 'in_tmem_layouts'}

try:
  ...call _in_attr_for_operand...
except ValueError as e:
  if 'Unknown attribute' in str(e):
    raise RuntimeError(f'Unsupported attribute name; check spelling/version: {e}') from e
  raise

When it happens

Trigger: Calling _in_attr_for_operand (or a public wrapper passing through an attr_name) with a name not in the supported set, e.g. 'out_tmem_layouts' before it was added, or a typo like 'in_layout'.

Common situations: Version drift where newer code references an attribute the installed inference_utils doesn't handle yet; copy-pasted code with a misspelled attribute name; extending the utilities with a new layout kind without updating the dispatch.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• Can only store to references (got {x_ref}).
• Can only infer one dimension
• Unsupported scope: {scope}
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• Unknown resize method "{s}"

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