{"record":{"id":"213955a17e804f96","repo":"jax-ml/jax","slug":"numpy-masked-arrays-are-not-supported-as-direct-in-213955","errorCode":null,"errorMessage":"numpy masked arrays are not supported as direct inputs to JAX functions. Use arr.filled() to convert the value to a standard numpy array.","messagePattern":"numpy masked arrays are not supported as direct inputs to JAX functions\\. Use arr\\.filled\\(\\) to convert the value to a standard numpy array\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/mlir.py","lineNumber":335,"sourceCode":"      return c_val\n  for t in type(val).__mro__:\n    handler = _constant_handlers.get(t)\n    if handler:\n      out = handler(val, aval)\n      assert _is_ir_values(out), (type(val), out)\n      return out\n  m = getattr(val, '__jax_array__', None)\n  if m is not None:\n    return ir_constant(m())\n  raise TypeError(f\"No constant handler for type: {type(val)}\")\n\n\ndef _numpy_array_constant(x: np.ndarray | np.generic) -> ir.Value:\n  return hlo.constant(_numpy_array_attribute(x))\n\n\ndef _masked_array_constant_handler(*args, **kwargs):\n  raise ValueError(\"numpy masked arrays are not supported as direct inputs to JAX functions. \"\n                   \"Use arr.filled() to convert the value to a standard numpy array.\")\n\nregister_constant_handler(np.ma.MaskedArray, _masked_array_constant_handler)\n\ndef _shape_dtype_struct_constant_handler(*args, **kwargs):\n  raise TypeError(\"A ShapeDtypeStruct does not have a value and cannot be \"\n                  \"used as a constant in a JAX function.\")\n\nregister_constant_handler(core.ShapeDtypeStruct,\n                          _shape_dtype_struct_constant_handler)\n\ndef _ndarray_constant_handler(val: np.ndarray | np.generic,\n                              aval: core.AbstractValue | None) -> IrValues:\n  \"\"\"Constant handler for ndarray literals, handling zero-size strides.\n\n  In most cases this function calls _numpy_array_constant(val) except it has\n  special handling of arrays with any strides of size zero: for those, it\n  generates appropriate calls to NumpyArrayConstant, Broadcast, and Transpose","sourceCodeStart":317,"sourceCodeEnd":353,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/mlir.py#L317-L353","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nresult = jitted_fn(masked_array)  # ValueError\n\n# after\nresult = jitted_fn(masked_array.filled(np.nan))","handlingStrategy":"validation","validationCode":"import numpy as np\n\ndef sanitize(arr):\n    if isinstance(arr, np.ma.MaskedArray):\n        return arr.filled(np.nan)\n    return arr\n\nx = sanitize(x)\njitted_fn(x)","typeGuard":"import numpy as np\n\ndef is_masked_array(a) -> bool:\n    return isinstance(a, np.ma.MaskedArray)","tryCatchPattern":"try:\n    out = jitted_fn(x)\nexcept ValueError as e:\n    if 'masked arrays' in str(e):\n        out = jitted_fn(x.filled(np.nan))\n    else:\n        raise","preventionTips":["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"],"tags":["numpy","jax","masked-array","input-validation"],"backgroundTag":"unsupported-input-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}