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

JAX explicitly rejects numpy masked arrays (np.ma.MaskedArray) as inputs because masked arrays carry a fill/mask semantic that has no equivalent in XLA/MLIR constants. A dedicated handler raises immediately so that silent wrong results (dropping the mask) cannot occur.

Source

Thrown at jax/_src/interpreters/mlir.py:335

      return c_val
  for t in type(val).__mro__:
    handler = _constant_handlers.get(t)
    if handler:
      out = handler(val, aval)
      assert _is_ir_values(out), (type(val), out)
      return out
  m = getattr(val, '__jax_array__', None)
  if m is not None:
    return ir_constant(m())
  raise TypeError(f"No constant handler for type: {type(val)}")


def _numpy_array_constant(x: np.ndarray | np.generic) -> ir.Value:
  return hlo.constant(_numpy_array_attribute(x))


def _masked_array_constant_handler(*args, **kwargs):
  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.")

register_constant_handler(np.ma.MaskedArray, _masked_array_constant_handler)

def _shape_dtype_struct_constant_handler(*args, **kwargs):
  raise TypeError("A ShapeDtypeStruct does not have a value and cannot be "
                  "used as a constant in a JAX function.")

register_constant_handler(core.ShapeDtypeStruct,
                          _shape_dtype_struct_constant_handler)

def _ndarray_constant_handler(val: np.ndarray | np.generic,
                              aval: core.AbstractValue | None) -> IrValues:
  """Constant handler for ndarray literals, handling zero-size strides.

  In most cases this function calls _numpy_array_constant(val) except it has
  special handling of arrays with any strides of size zero: for those, it
  generates appropriate calls to NumpyArrayConstant, Broadcast, and Transpose

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Call arr.filled() (optionally with an explicit fill value like arr.filled(np.nan) or 0) before passing to JAX
Convert with np.asarray(arr) to drop the mask if the mask is irrelevant
Fix upstream loading (e.g. xarray open_dataset(decode_cf=...) settings) to not produce masked arrays

Example fix

# before
result = jitted_fn(masked_array)  # ValueError

# after
result = jitted_fn(masked_array.filled(np.nan))

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np

def sanitize(arr):
    if isinstance(arr, np.ma.MaskedArray):
        return arr.filled(np.nan)
    return arr

x = sanitize(x)
jitted_fn(x)

Type guard

import numpy as np

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

Try / catch

try:
    out = jitted_fn(x)
except ValueError as e:
    if 'masked arrays' in str(e):
        out = jitted_fn(x.filled(np.nan))
    else:
        raise

Prevention

Sanitize loaded data at ingestion: arr.filled(...) right after netCDF/xarray reads
Assert inputs are plain ndarrays in a thin wrapper around jitted entry points

When it happens

Trigger: Passing an np.ma.MaskedArray directly as an argument to a jitted JAX function, or as a constant captured in a closure during tracing. Common after data loading pipelines (e.g. netCDF, climate data via xarray with _FillValue handling) that produce masked arrays.

Common situations: Scientific data pipelines reading netCDF/HDF5 with missing values, np.genfromtxt with missing_values=True, or arithmetic that returns masked arrays; version changes where previously masked arrays were silently coerced.

Related errors

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