jax-ml/jax · error · ValueError

numpy masked arrays are not supported as direct inputs to JA

Error message

numpy masked arrays are not supported as direct inputs to JAX functions. Use arr.filled() to convert the value to a standard numpy array.

What it means

Numpy masked arrays (np.ma.MaskedArray) carry a mask that JAX cannot represent, so when one is passed as an argument to a sharded/jitted computation, pxla deliberately raises instead of silently dropping the mask. The fix is to materialize the mask via arr.filled().

Source

Thrown at jax/_src/interpreters/pxla.py:174

  d = sharding._device_assignment[0]
  shard_shape = sharding.shard_shape(aval.shape)
  try:
    # TODO(yashkatariya): Replace this with normal `==` check once CPU supports
    # int4.
    return is_user_xla_layout_equal(
        curr_layout,
        Layout.from_pjrt_layout(
            d.client.get_default_layout(aval.dtype, shard_shape, d)))
  except _jax.JaxRuntimeError as e:
    msg, *_ = e.args
    if isinstance(msg, str) and msg.startswith("UNIMPLEMENTED"):
      return True
    else:
      raise


def _masked_array_error(xs, shardings, layouts, copy_semantics):
  raise ValueError("numpy masked arrays are not supported as direct inputs to JAX functions. "
                   "Use arr.filled() to convert the value to a standard numpy array.")
shard_arg_handlers[np.ma.MaskedArray] = _masked_array_error

def _shard_np_array(xs, shardings, layouts, copy_semantics):
  results = []
  batch_xs, batch_cs, batch_shardings, batch_indices = [], [], [], []
  for i, (x, sharding, layout, cs) in enumerate(
      zip(xs, shardings, layouts, copy_semantics)):
    if x.dtype == dtypes.float0:
      x = np.zeros(x.shape, dtype=np.dtype(bool))
    if layout is not None:
      results.append(api.device_put(x, Format(layout, sharding)))
    else:
      if config.use_cpp_shard_args.value:
        results.append(None)
        batch_xs.append(x)  # Accumulate arguments to `_jax.shard_args`
        batch_shardings.append(sharding)
        batch_indices.append(i)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert before the call: use arr.filled(fill_value) (e.g. filled(np.nan) or 0) to get a plain ndarray
2. Alternatively use np.ma.getdata(arr) plus manual mask handling as a separate array input
3. If using xarray, convert with .values or .fillna first

Example fix

# before
masked = np.ma.masked_invalid(data)
out = jitted_fn(masked)

# after
masked = np.ma.masked_invalid(data)
out = jitted_fn(masked.filled(np.nan))

def to_jax_compatible(arr):
    if isinstance(arr, np.ma.MaskedArray):
        return arr.filled(np.nan)
    return arr
args = jax.tree.map(to_jax_compatible, args)

import numpy as np

def is_masked_array(x) -> bool:
    return isinstance(x, np.ma.MaskedArray)

When it happens

Trigger: Passing an np.ma.MaskedArray (e.g. from np.ma.masked_invalid, masked_where, or netCDF-style data loaders) as an argument to a jitted/pjit/sharded function. JAX registers an explicit shard_arg_handler for np.ma.MaskedArray that always raises.

Common situations: Scientific data pipelines (climate/geo data via xarray/netCDF often yields masked arrays); NaN handling with np.ma.masked_invalid; forgetting a .filled() after masked preprocessing.

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• 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.

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