Repository: google/jax
Issue: JAX silently discards masks in numpy masked arrays

Reported in https://github.com/data-apis/array-api/issues/582
```python
>>> jnp.sin(np.ma.array([1, 2, 3], mask=[True, False, True]))
Array([0.84147096, 0.9092974 , 0.14112   ], dtype=float32)
```

Changed files:
jax/interpreters/xla.py, tests/jax_jit_test.py, jax/interpreters/mlir.py, tests/api_test.py, tests/lax_numpy_test.py, jax/interpreters/pxla.py, jax/_src/abstract_arrays.py, jax/_src/dispatch.py

diff --git a/jax/_src/abstract_arrays.py b/jax/_src/abstract_arrays.py
index db47ee369..588e57cf3 100644
--- a/jax/_src/abstract_arrays.py
+++ b/jax/_src/abstract_arrays.py
@@ -51,6 +51,13 @@ def canonical_concrete_aval(val, weak_type=None):
   return ConcreteArray(dtypes.canonicalize_dtype(np.result_type(val)), val,
                        weak_type=weak_type)
 
+def masked_array_error(*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.")
+
+core.pytype_aval_mappings[np.ma.MaskedArray] = masked_array_error
+ad_util.jaxval_zeros_likers[np.ma.MaskedArray] = masked_array_error
+
 for t in array_types:
   core.pytype_aval_mappings[t] = canonical_concrete_aval
   ad_util.jaxval_zeros_likers[t] = zeros_like_array
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
index 8372e4a60..b67c158b2 100644
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -1275,6 +1275,10 @@ def device_put(x, device: Optional[Device] = None) -> Tuple[Any, ...]:
 # TODO(phawkins): update users.
 xla.device_put = device_put
 
+def _device_put_masked_array(x, device: Optional[Device]):
+  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.")
+
 def _device_put_array(x, device: Optional[Device]):
   backend = xb.get_device_backend(device)
   if x.dtype == dtypes.float0:
@@ -1293,6 +1297,7 @@ _scalar_types = dtypes.python_scalar_dtypes.keys()
 
 device_put_handlers: Dict[Any, Callable[[Any, Optional[Device]],
                                         Tuple[Any, ...]]] = {}
+device_put_handlers[np.ma.MaskedArray] = _device_put_masked_array
 device_put_handlers.update((t, _device_put_array) for t in array_types)
 device_put_handlers.update((t, _device_put_scalar) for t in _scalar_types)
 device_put_handlers[core.Token] = _device_put_token
diff --git a/jax/interpreters/mlir.py b/jax/interpreters/mlir.py
index 3f273fc09..f010c8561 100644
--- a/jax/interpreters/mlir.py
+++ b/jax/interpreters/mlir.py
@@ -241,6 +241,11 @@ def _numpy_array_constant(x: np.ndarray, canonicalize_types
   return (hlo.ConstantOp(attr).result,)
 
 
+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 _ndarray_constant_handler(val: np.ndarray, canonicalize_types
                              ) -> Sequence[ir.Value]:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
index 639b848ac..02d8b2bda 100644
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -421,6 +421,11 @@ def _shard_token(x, devices, indices):
   return device_put(np.zeros((), dtype=np.dtype(np.bool_)), devices, replicate=True)
 shard_arg_handlers[core.Token] = _shard_token
 
+def _masked_array_error(*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.")
+shard_arg_handlers[np.ma.MaskedArray] = _masked_array_error
+
 def _shard_array(x, devices, indices):
   if x.dtype == dtypes.float0:
     x = np.zeros(x.shape, dtype=np.dtype(bool))
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 268f0dbe5..d81bc830f 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -243,6 +243,9 @@ def canonicalize_dtype(x):
     return canonicalize_dtype(x.__jax_array__())
   raise TypeError(f"No canonicalize_dtype handler for type: {type(x)}")
 
+def _canonicalize_masked_array_dtype(x):
+  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.")
 
 def _canonicalize_ndarray_dtype(x):
   return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
@@ -260,6 +263,7 @@ for t in device_array.device_array_types:
 canonicalize_dtype_handlers.update(
     (t, _canonicalize_ndarray_dtype) for t in numpy_scalar_types)
 canonicalize_dtype_handlers[np.ndarray] = _canonicalize_ndarray_dtype
+canonicalize_dtype_handlers[np.ma.MaskedArray] = _canonicalize_masked_array_dtype
 canonicalize_dtype_handlers.update(
     (t, partial(_canonicalize_python_scalar_dtype, t)) for t in _scalar_types)
 canonicalize_dtype_handlers[core.Token] = identity
diff --git a/tests/api_test.py b/tests/api_test.py
index 317e4f5a0..b2c0de268 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -589,6 +589,12 @@ class CPPJitTest(jtu.BufferDonationTestCase):
         ".* '.*int32.*' of type <.*_ScalarMeta.*> is not a valid JAX type",
         lambda: self.jit(f)(jnp.int32))
 
+  def test_jit_masked_array(self):
+    x = np.ma.array([1, 2, 3], mask=[True, False, True])
+    f = self.jit(lambda x: x)
+    with self.assertRaisesRegex(ValueError, "numpy masked arrays are not supported"):
+      f(x)
+
   def test_jit_on_all_devices(self):
     # Verifies we can run the same computation on every device present, even
     # if they are, for example, different models of GPU.
diff --git a/tests/jax_jit_test.py b/tests/jax_jit_test.py
index 7f6bc35a7..5f0af752d 100644
--- a/tests/jax_jit_test.py
+++ b/tests/jax_jit_test.py
@@ -39,6 +39,16 @@ class JaxJitTest(jtu.JaxTestCase):
         jaxlib.jax_jit._is_float0(np.zeros((5, 5), dtype=jax.float0)))
     self.assertFalse(jaxlib.jax_jit._is_float0(np.zeros((5, 5))))
 
+  @parameterized.parameters([jax.device_put, _cpp_device_put])
+  def test_device_put_on_numpy_masked_array(self, device_put_function):
+    # TODO(jakevdp): add appropriate logic to jaxlib device_put and update this test.
+    if device_put_function is _cpp_device_put:
+      self.skipTest("cpp device_put does not yet reject masked arrays.")
+    device = jax.devices()[0]
+    value = np.ma.array([1, 2, 3], mask=[True, False, True])
+    with self.assertRaisesRegex(ValueError, "numpy masked arrays are not supported"):
+      device_put_function(value, device=device)
+
   @parameterized.parameters([jax.device_put, _cpp_device_put])
   def test_device_put_on_numpy_scalars(self, device_put_function):
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index bb194f54b..678562b71 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2904,6 +2904,14 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  def testArrayFromMasked(self):
+    args_maker = lambda: [np.ma.array([1, 2], mask=[True, False])]
+    # Like np.array, jnp.array strips the mask from masked array inputs.
+    self._CheckAgainstNumpy(np.array, jnp.array, args_maker)
+    # Under JIT, masked arrays are flagged as invalid.
+    with self.assertRaisesRegex(ValueError, "numpy masked arrays are not supported"):
+      jax.jit(jnp.asarray)(*args_maker())
+
   @jtu.sample_product(
     [dict(arg=arg, dtype=dtype, ndmin=ndmin)
       for arg, dtypes in [