jax-ml/jax · error · TypeError

Argument to to_dlpack must be a jax.Array, got {type(x)}

Error message

Argument to to_dlpack must be a jax.Array, got {type(x)}

What it means

jax.dlpack.to_dlpack only accepts concrete jax.Array (ArrayImpl) instances; anything else (tracers, numpy arrays, torch tensors, Python scalars) is rejected with a TypeError.

Source

Thrown at jax/_src/dlpack.py:132

    copy: a boolean indicating whether or not to copy the input. If
      ``copy=True`` then the function must always copy. When
      ``copy=False`` then the function must never copy, and must raise an error
      when a copy is deemed necessary. If ``copy=None`` then the function must
      avoid a copy if possible but may copy if needed.

  Returns:
    A DLPack PyCapsule object.

  Note:
    While JAX arrays are always immutable, ``DLPackManagedTensor`` buffers
    cannot be marked as immutable, and it is possible for processes external
    to JAX to mutate them in-place. If a DLPack buffer derived from a JAX array
    is mutated, it may lead to undefined behavior when using the associated JAX
    array. When JAX eventually supports ``DLManagedTensorVersioned``
    (DLPack 1.0), it will be possible to specify that a buffer is read-only.
  """
  if not isinstance(x, array.ArrayImpl):
    raise TypeError("Argument to to_dlpack must be a jax.Array, "
                    f"got {type(x)}")

  device = None
  dl_device_type, local_hardware_id = dl_device if dl_device else (None, None)
  if dl_device_type:
    try:
      dl_device_platform = _DL_DEVICE_TO_PLATFORM[dl_device_type]
      backend = xla_bridge.get_backend(dl_device_platform)
      device = backend.device_from_local_hardware_id(local_hardware_id)
    except KeyError:
      # https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html
      # recommends using BufferError.
      raise BufferError(
          "The device specification passed to to_dlpack contains an"
          f" unsupported device type (DLDeviceType: {dl_device_type})"
      ) from None

  # As new versions are adopted over time, we can maintain some legacy paths

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert to a JAX array first: `x = jax.numpy.asarray(x)`
Check `isinstance(x, jax.Array)` before calling to_dlpack
For numpy inputs, use torch/other dlpack producers or convert via device_put

Example fix

# before
jax.dlpack.to_dlpack(np_array)

# after
jax.dlpack.to_dlpack(jax.numpy.asarray(np_array))

Defensive patterns

Strategy: type-guard

Validate before calling

import jax
if not isinstance(x, jax.Array):
    x = jax.numpy.asarray(x)

Type guard

import jax
def is_jax_array(x) -> bool:
    return isinstance(x, jax.Array)

Try / catch

try:
    jax.dlpack.to_dlpack(x)
except TypeError:
    jax.dlpack.to_dlpack(jax.numpy.asarray(x))

Prevention

Normalize inputs with jnp.asarray at API boundaries
Never call to_dlpack inside jit on tracers

When it happens

Trigger: Passing a NumPy array, torch.Tensor, Python scalar, or a JAX tracer (inside jit) to jax.dlpack.to_dlpack().

Common situations: Interop helpers that receive 'array-like' inputs; passing a jitted function's internal tracer; passing numpy arrays after conversion assumptions break.

Related errors

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