jax-ml/jax · error · ValueError

`{device_or_sharding = }` was passed to`canonicalize_or_get_

Error message

`{device_or_sharding = }` was passed to`canonicalize_or_get_default_platform`, only xc.Device, Sharding, None or str values are supported.

What it means

_get_platform (used by asarray/array to resolve the target device) accepts only an xc.Device, a Sharding, None, or a platform string like 'cpu'/'gpu'/'tpu'. Any other type — int device id, torch device, etc. — cannot be canonicalized and raises ValueError.

Source

Thrown at jax/_src/numpy/array_constructors.py:360

  return out_array


def _get_platform(
    device_or_sharding: xc.Device | Sharding | None | str) -> str:
  """Get device_or_sharding platform or look up config.default_device.value."""
  if isinstance(device_or_sharding, xc.Device):
    return device_or_sharding.platform
  elif isinstance(device_or_sharding, Sharding):
    return list(device_or_sharding.device_set)[0].platform
  elif isinstance(device_or_sharding, str):
    return device_or_sharding
  elif device_or_sharding is None:
    if config.default_device.value is None:
      return xla_bridge.default_backend()
    else:
      return _get_platform(config.default_device.value)
  else:
    raise ValueError(f"`{device_or_sharding = }` was passed to"
                     "`canonicalize_or_get_default_platform`, only xc.Device,"
                     " Sharding, None or str values are supported.")


def _convert_to_array_if_dtype_fails(x: ArrayLike) -> ArrayLike:
  try:
    dtypes.dtype(x)
  except TypeError:
    return np.asarray(x)
  else:
    return x


@export
def asarray(a: Any, dtype: DTypeLike | None = None, order: str | None = None,
            *, copy: bool | None = None,
            device: xc.Device | Sharding | None = None,
            out_sharding: NamedSharding | P | None = None) -> Array:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.devices('gpu')[0] to get a Device object, or the string platform name
2. For sharded output, pass a jax.sharding.Sharding instance
3. Wrap integer ids: jax.devices()[device_id]

Example fix

# before
a = jnp.asarray(x, device=0)
# after
import jax
a = jnp.asarray(x, device=jax.devices('gpu')[0])

import jax
from jax.sharding import Sharding
from jaxlib.xla_extension import Device
def valid_device(d):
    return d is None or isinstance(d, (Device, Sharding, str))

def is_valid_jax_device(d) -> bool:
    import jax
    from jax.sharding import Sharding
    return d is None or isinstance(d, (jax.Device, Sharding, str))

null

When it happens

Trigger: jnp.asarray(x, device=0) or jnp.asarray(x, device=torch.device('cuda')) — a device argument that is not Device/Sharding/None/str.

Common situations: Porting PyTorch or NumPy-CUDA code that uses integer device IDs; passing jax.local_devices(1)[0]'s id instead of the Device object.

Related errors

• Argument '{arg}' of type {type(arg)} is not a valid JAX type
• SymbolicScope constraint must be a string: got {repr(c_str)}
• Custom VJP bwd rule {self.bwd} was registered with defvjp_wi
• All objects to concatenate must be arrays, got {}.
• operand to reduce_window_sum must have a number dtype, got {

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