jax-ml/jax · error · TypeError

{op} has {name} of an unexpected type: {result}

Error message

{op} has {name} of an unexpected type: {result}

What it means

_array_attr requires the named attribute (in_layouts, out_layouts, transforms, tmem layouts) to be an MLIR ir.ArrayAttr. If the attribute exists but has another type (e.g. a single layout attribute, a string, or a malformed encoding), a TypeError is raised showing the unexpected value. This usually indicates the op was serialized/parsed or constructed incorrectly.

Source

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

  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. Rebuild or re-attach the attribute as an ir.ArrayAttr, e.g. op.attributes['in_layouts'] = ir.ArrayAttr.get([layout_attr])
2. Inspect the printed attribute (the error message includes {result}) to identify why it is not an array and fix the construction site
3. If parsing textual MLIR, ensure array attributes use array syntax and the correct dialect types

Example fix

# before
op.attributes['in_layouts'] = some_layout_attr  # not an ArrayAttr

# after
op.attributes['in_layouts'] = ir.ArrayAttr.get([some_layout_attr])

import iree.compiler.dialects.ir as ir

attr = op.attributes.get('in_layouts')
if attr is not None and not isinstance(attr, ir.ArrayAttr):
  op.attributes['in_layouts'] = ir.ArrayAttr.get([attr])

def is_array_attr(op, name) -> bool:
  import iree.compiler.dialects.ir as ir
  return isinstance(op.attributes.get(name), ir.ArrayAttr)

try:
  layouts = inference_utils.in_layouts(op)
except TypeError as e:
  if 'unexpected type' in str(e):
    op.attributes['in_layouts'] = ir.ArrayAttr.get([op.attributes['in_layouts']])
    layouts = inference_utils.in_layouts(op)
  else:
    raise

When it happens

Trigger: Calling in_layouts/out_layouts/in_transforms/out_transforms/in_tmem_layouts/out_tmem_layouts on an op whose attribute of that name is present but not an ir.ArrayAttr, e.g. a single layout object stored directly instead of a one-element array, or an attribute decoded from a mismatched MLIR context/version.

Common situations: Round-tripping MLIR through textual IR where an attribute got attached without array syntax ([...]); building ops with a helper that sets a bare attribute; version mismatch between the MLIR bindings and the Mosaic dialect definitions.

Related errors

• {op} does not have an {name} attribute
• Expected an index-typed index
• Only support memref.cast where the input and output types ar
• Only unit strides are supported but got {op.static_strides}.
• Transposed memrefs are not supported in ExpandShapeOp.

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