jax-ml/jax · error · ValueError

Can't instantiate {self} with arguments.

Error message

Can't instantiate {self} with arguments.

What it means

Raised by Layout.to_mgpu() on a plain (already concrete) Layout wrapper: no construction arguments are accepted because the layout needs no parameters. Passing any positional or keyword args to to_mgpu for a parameterless layout is a usage error.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1825

    return ReducedLayout(self, axes)

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    raise NotImplementedError


@dataclasses.dataclass(frozen=True)
class ParameterizedLayout(SomeLayout):
  layout_cls: Layout | TMEMLayout
  args: Sequence[Any]
  kwargs: Any

  def __post_init__(self):
    object.__setattr__(self, "args", tuple(self.args))
    object.__setattr__(self, "kwargs", frozen_dict.FrozenDict(self.kwargs))

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    if args or kwargs:
      raise ValueError(f"Can't instantiate {self} with arguments.")
    return self.layout_cls.to_mgpu(*self.args, **self.kwargs)


@dataclasses.dataclass(frozen=True)
class ReducedLayout(SomeLayout):
  layout: SomeLayout
  axes: Sequence[int]

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    if args or kwargs:
      raise ValueError(f"Can't instantiate {self} with arguments.")
    layout = self.layout.to_mgpu()
    if not isinstance(layout, mgpu.TiledLayout):
      raise ValueError("Only TiledLayout supports reductions.")
    return layout.reduce(self.axes)


class Layout(SomeLayout, enum.Enum):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the arguments: call layout.to_mgpu()
2. If parameters are genuinely needed, use a ParameterizedLayout created via the __call__ operator instead of a plain layout

Example fix

// before
layout.to_mgpu(128)

// after
layout.to_mgpu()

if args or kwargs:
    raise TypeError('concrete layouts take no to_mgpu arguments')
layout.to_mgpu()

def call_to_mgpu(layout, *args, **kwargs):
    from jax._src.pallas.mosaic_gpu.core import ParameterizedLayout
    if not isinstance(layout, ParameterizedLayout) and (args or kwargs):
        raise TypeError('parameterless layout: call to_mgpu() with no args')
    return layout.to_mgpu(*args, **kwargs)

try:
    mgpu_layout = layout.to_mgpu(*args, **kwargs)
except ValueError:
    mgpu_layout = layout.to_mgpu()

When it happens

Trigger: Calling `layout.to_mgpu(shape)` or `layout.to_mgpu(bits=128)` on a SomeLayout instance that wraps a concrete layout_cls taking no init arguments.

Common situations: Generic code that forwards user kwargs to to_mgpu for all layouts, including parameterless ones; refactoring from ParameterizedLayout (which does take args) to a fixed Layout without removing the call-site arguments.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• dims and idxs must have the same length
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128

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