jax-ml/jax · error · ValueError

Cannot specify both out_shape and out_type. Use out_type.

Error message

Cannot specify both out_shape and out_type. Use out_type.

What it means

In `jax.experimental.pallas.mosaic_gpu.kernel`, `out_shape` is the deprecated name for what is now `out_type`. You may pass one or the other, not both; supplying both raises this ValueError telling you to use `out_type`.

Source

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

        interpret=interpret,
        debug=debug,
        **mesh_kwargs,
    )

  if (
      not isinstance(out_shape, api.NotSpecified)
      or not isinstance(scratch_shapes, api.NotSpecified)
  ):
    deprecations.warn(
        "jax-pallas-mgpu-shapes-types",
        "The out_shape and scratch_shapes arguments to plgpu.kernel are"
        " deprecated. Use out_type and scratch_types instead.",
        stacklevel=2,
    )

  if not isinstance(out_shape, api.NotSpecified):
    if not isinstance(out_type, api.NotSpecified):
      raise ValueError(
          "Cannot specify both out_shape and out_type. Use out_type."
      )
    out_type = out_shape
  elif isinstance(out_type, api.NotSpecified):
    out_type = ()

  if not isinstance(scratch_shapes, api.NotSpecified):
    if not isinstance(scratch_types, api.NotSpecified):
      raise ValueError(
          "Cannot specify both scratch_shapes and scratch_types. Use"
          " scratch_types."
      )
    scratch_types = scratch_shapes
  elif isinstance(scratch_types, api.NotSpecified):
    scratch_types = ()

  if unwrap_out := not isinstance(out_type, (tuple, list)):
    out_type = (out_type,)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Delete `out_shape=` and keep only `out_type=` (they are equivalent, out_type is the modern name)
2. If you must support old JAX versions, branch on `jax.__version__` or use try/except TypeError
3. Silence the deprecation path entirely by never passing out_shape

Example fix

# before
kernel_fn = mgpu.kernel(pallas_fn, out_shape=(128, jnp.float32), out_type=(128, jnp.float32), grid=grid)
# after
kernel_fn = mgpu.kernel(pallas_fn, out_type=(128, jnp.float32), grid=grid)

import inspect
def make_kernel(fn, **kw):
    params = inspect.signature(mgpu.kernel).parameters
    if 'out_type' in kw and 'out_shape' in kw:
        del kw['out_shape']  # prefer modern name
    if 'out_type' not in params and 'out_type' in kw:
        kw['out_shape'] = kw.pop('out_type')
    return mgpu.kernel(fn, **kw)

When it happens

Trigger: Calling `mosaic_gpu.kernel(f, out_shape=..., out_type=..., ...)` with both keyword arguments non-NotSpecified.

Common situations: Migrating older Pallas code to a newer JAX where `out_type` was introduced; automated codemods that added `out_type` without removing the legacy `out_shape`; merging code from two branches that each used a different spelling.

Related errors

• Cannot specify both scratch_shapes and scratch_types. Use sc
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Unsupported dtype: {ref.dtype}
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp

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