{"record":{"id":"022385d77ec98229","repo":"jax-ml/jax","slug":"numpy-masked-arrays-are-not-supported-as-direct-in-022385","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":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/pxla.py","lineNumber":174,"sourceCode":"  d = sharding._device_assignment[0]\n  shard_shape = sharding.shard_shape(aval.shape)\n  try:\n    # TODO(yashkatariya): Replace this with normal `==` check once CPU supports\n    # int4.\n    return is_user_xla_layout_equal(\n        curr_layout,\n        Layout.from_pjrt_layout(\n            d.client.get_default_layout(aval.dtype, shard_shape, d)))\n  except _jax.JaxRuntimeError as e:\n    msg, *_ = e.args\n    if isinstance(msg, str) and msg.startswith(\"UNIMPLEMENTED\"):\n      return True\n    else:\n      raise\n\n\ndef _masked_array_error(xs, shardings, layouts, copy_semantics):\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.\")\nshard_arg_handlers[np.ma.MaskedArray] = _masked_array_error\n\ndef _shard_np_array(xs, shardings, layouts, copy_semantics):\n  results = []\n  batch_xs, batch_cs, batch_shardings, batch_indices = [], [], [], []\n  for i, (x, sharding, layout, cs) in enumerate(\n      zip(xs, shardings, layouts, copy_semantics)):\n    if x.dtype == dtypes.float0:\n      x = np.zeros(x.shape, dtype=np.dtype(bool))\n    if layout is not None:\n      results.append(api.device_put(x, Format(layout, sharding)))\n    else:\n      if config.use_cpp_shard_args.value:\n        results.append(None)\n        batch_xs.append(x)  # Accumulate arguments to `_jax.shard_args`\n        batch_shardings.append(sharding)\n        batch_indices.append(i)","sourceCodeStart":156,"sourceCodeEnd":192,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/pxla.py#L156-L192","documentation":"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().","triggerScenarios":"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.","commonSituations":"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.","solutions":["Convert before the call: use arr.filled(fill_value) (e.g. filled(np.nan) or 0) to get a plain ndarray","Alternatively use np.ma.getdata(arr) plus manual mask handling as a separate array input","If using xarray, convert with .values or .fillna first"],"exampleFix":"# before\nmasked = np.ma.masked_invalid(data)\nout = jitted_fn(masked)\n\n# after\nmasked = np.ma.masked_invalid(data)\nout = jitted_fn(masked.filled(np.nan))","handlingStrategy":"validation","validationCode":"def to_jax_compatible(arr):\n    if isinstance(arr, np.ma.MaskedArray):\n        return arr.filled(np.nan)\n    return arr\nargs = jax.tree.map(to_jax_compatible, args)","typeGuard":"import numpy as np\n\ndef is_masked_array(x) -> bool:\n    return isinstance(x, np.ma.MaskedArray)","tryCatchPattern":null,"preventionTips":["Call .filled() on any np.ma array before passing to JAX","Watch for masked arrays emerging from netCDF/xarray pipelines"],"tags":["jax","numpy","masked-array","input-validation","sharding"],"backgroundTag":"unsupported-input-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}