jax-ml/jax · error · TypeError

Invalid value received for the sharding argument. Expected v

Error message

Invalid value received for the sharding argument. Expected values are `None` or an instance of `jax.Sharding`. Got {sharding} of type {type(sharding)}

What it means

Raised by the same layout module constructor when the `sharding` argument is neither None nor a jax.Sharding instance. JAX requires sharding specs to be real Sharding objects (NamedSharding, SingleDeviceSharding, etc.) so it can compute device assignments. Any other value, including duck-typed or pickled objects from mismatched versions, fails this isinstance gate.

Source

Thrown at jax/_src/layout.py:154

    # If layout is concrete and sharding is not, error.
    if isinstance(layout, Layout) and sharding is None:
      raise ValueError(
          'Sharding has to be concrete when layout is of type'
          f' {type(layout)}. Please pass a'
          ' `jax.sharding.NamedSharding` or'
          ' `jax.sharding.SingleDeviceSharding` to the sharding argument. Got'
          f' sharding {sharding}'
      )
    if not isinstance(
        layout, (Layout, type(None), AutoLayoutSingleton)):
      raise TypeError(
          'Invalid value received for the layout argument.'
          ' Expected values are `None`, `Layout.AUTO` or an'
          f' instance of `Layout`. Got {layout} of'
          f' type {type(layout)}'
      )
    if not isinstance(sharding, (Sharding, type(None))):
      raise TypeError(
          'Invalid value received for the sharding argument. Expected values'
          ' are `None` or an instance of `jax.Sharding`. Got'
          f' {sharding} of type {type(sharding)}')

    self.layout = layout
    self.sharding = sharding

  def __repr__(self):
    return f'Format(layout={self.layout}, sharding={self.sharding})'

  def __hash__(self):
    return hash((self.layout, self.sharding))

  def __eq__(self, other):
    if not isinstance(other, Format):
      return False
    return (self.layout == other.layout and
            self.sharding == other.sharding)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a jax.sharding.* instance such as NamedSharding(mesh, P('x')) or None
2. Verify only one jax installation is imported (check jax.__file__ on both producer and consumer)
3. Recreate the sharding object in the current process instead of unpickling from another version
4. Double-check you haven't swapped layout and sharding keyword arguments

Example fix

# before
sharding = ('data',)  # tuple, not a Sharding

# after
import jax
sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('data'))

import jax
ok = sharding is None or isinstance(sharding, jax.sharding.Sharding)

import jax
def is_valid_sharding(x) -> bool:
    return x is None or isinstance(x, jax.sharding.Sharding)

try:
    f = jax.jit(fun, in_shardings=sharding)
except TypeError as e:
    if 'sharding argument' in str(e):
        f = jax.jit(fun)  # fall back to default sharding
    else: raise

When it happens

Trigger: Passing a string ('sharding'), a partition spec tuple, or a numpy array as `sharding`; passing a Sharding object created in a different jax version/install whose class identity differs; mixing up argument order with layout.

Common situations: Argument mix-ups in jit(..., in_shardings=...) style calls; loading pickled sharding objects across JAX versions; having two jax installs (e.g. pip + conda) so isinstance fails across module identities.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• to_dlpack can only pack a dlpack tensor from an array on a s
• __dlpack__ only supported for unsharded arrays.
• unsafe_buffer_pointer() is supported only for unsharded arra

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