Repository: google-deepmind/penzai

Issue:
[BUG] Order sensitivity for implicit casting of jax Arrays with infix operators
Adding shapeless Jax arrays to named arrays with infix notation works from the right, but not the left. This does not occur when you nmap the `jnp.add` function directly, so a potential solution might be to replace the `NamedArrayBase` dunder methods with the `jax.numpy` version. e.g. 

``` __add__ = _nmap_with_doc(operator.add, "jax.Array.__add__")``` -> 
``` __add__ = _nmap_with_doc(jax.numpy.add,"jax.Array.__add__")```

Reproducable example: 

```python
# Broken
import jax.numpy as jnp
from penzai import pz

def test_right_const(arr):
    return arr + jnp.array(3.) 

def test_left_const(arr):
    return jnp.array(3.) + arr     

arr = pz.nx.arange('arr', 3)

test_right_const(arr) # works
test_left_const(arr) # fails
```

```
ValueError: Only NamedArray(View)s with no named axes can be converted to JAX arrays. Use `unwrap` or `untag` to assign positions to named axes first, or use `penzai.named_axes.nmap` with a JAX function instead.
```

```python
# Working
def no_infix_right_const(arr):
    return pz.nx.nmap(jnp.add)(arr, jnp.array(3.))

def no_infix_left_const(arr):
    return pz.nx.nmap(jnp.add)(jnp.array(3.), arr)

# these both work
no_infix_left_const(arr)
no_infix_right_const(arr)
```

Discussion:

Patch:
diff --git a/penzai/core/named_axes.py b/penzai/core/named_axes.py
index 2e74ef4..7974978 100644
--- a/penzai/core/named_axes.py
+++ b/penzai/core/named_axes.py
@@ -867,29 +867,6 @@ class NamedArrayBase(abc.ABC):
           index_thunks.append(_StaticThunk(c))
       return _jitted_nmapped_getitem(self, tuple(index_thunks))
 
-  # Array conversion operators
-  def __array__(self, *args, **kwargs):
-    """Converts a named array with no named axes to a Numpy array."""
-    if self.named_shape:
-      raise ValueError(
-          "Only NamedArray(View)s with no named axes can be converted to numpy"
-          " arrays. Use `unwrap` or `untag` to assign positions to named axes"
-          " first, or use `penzai.named_axes.nmap` with a JAX function instead."
-      )
-    else:
-      return np.array(self.unwrap(), *args, **kwargs)
-
-  def __jax_array__(self):
-    """Converts a named array with no named axes to a JAX array."""
-    if self.named_shape:
-      raise ValueError(
-          "Only NamedArray(View)s with no named axes can be converted to JAX"
-          " arrays. Use `unwrap` or `untag` to assign positions to named axes"
-          " first, or use `penzai.named_axes.nmap` with a JAX function instead."
-      )
-    else:
-      return self.unwrap()
-
   # Iteration. Note that we *must* implement this to avoid Python simply trying
   # to run __getitem__ until it raises IndexError, because we won't raise
   # IndexError (since JAX clips array indices).


Test patch:
diff --git a/tests/named_axes_test.py b/tests/named_axes_test.py
index ca99da7..d43910e 100644
--- a/tests/named_axes_test.py
+++ b/tests/named_axes_test.py
@@ -569,27 +569,6 @@ class NamedAxesTest(parameterized.TestCase):
           ),
       )
 
-    with self.subTest("binop_scalar"):
-      chex.assert_trees_all_equal(
-          array_a - 100,
-          named_axes.NamedArray(
-              named_axes=collections.OrderedDict({"foo": 2, "bar": 3}),
-              data_array=np.array(
-                  [[-100, -99, -98], [-97, -96, -95]],
-                  dtype=np.dtype("int32"),
-              ),
-          ),
-      )
-      chex.assert_trees_all_equal(
-          100 - array_a,
-          named_axes.NamedArray(
-              named_axes=collections.OrderedDict({"foo": 2, "bar": 3}),
-              data_array=np.array(
-                  [[100, 99, 98], [97, 96, 95]], dtype=np.dtype("int32")
-              ),
-          ),
-      )
-
     with self.subTest("array_method"):
       # Summing over no axes does nothing
       chex.assert_trees_all_equal(array_a.sum(), array_a)
@@ -625,6 +604,43 @@ class NamedAxesTest(parameterized.TestCase):
       self.assertEqual(int(array_scalar), 3)
       self.assertTrue(bool(array_scalar))
 
+  @parameterized.parameters("python", "numpy", "jax")
+  def test_binop_lifts(self, which):
+    named_array = named_axes.wrap(jnp.arange(2 * 3).reshape((2, 3))).tag(
+        "foo", "bar"
+    )
+    if which == "python":
+      other = 100
+    elif which == "numpy":
+      other = np.array(100)
+    elif which == "jax":
+      other = jnp.array(100)
+    else:
+      raise ValueError(which)
+
+    with self.subTest("left"):
+      chex.assert_trees_all_equal(
+          named_array - other,
+          named_axes.NamedArray(
+              named_axes=collections.OrderedDict({"foo": 2, "bar": 3}),
+              data_array=np.array(
+                  [[-100, -99, -98], [-97, -96, -95]],
+                  dtype=np.dtype("int32"),
+              ),
+          ),
+      )
+
+    with self.subTest("right"):
+      chex.assert_trees_all_equal(
+          other - named_array,
+          named_axes.NamedArray(
+              named_axes=collections.OrderedDict({"foo": 2, "bar": 3}),
+              data_array=np.array(
+                  [[100, 99, 98], [97, 96, 95]], dtype=np.dtype("int32")
+              ),
+          ),
+      )
+
   def test_named_array_equality(self):
     array_a = named_axes.wrap(jnp.arange(2 * 3).reshape((2, 3))).tag(
         "foo", "bar"
@@ -636,37 +652,6 @@ class NamedAxesTest(parameterized.TestCase):
     result = array_a == array_b
     self.assertEqual(result.named_shape, {"foo": 2, "bar": 3, "baz": 5})
 
-  def test_named_array_conversion(self):
-    array_a = named_axes.wrap(jnp.arange(6))
-
-    with self.subTest("convert_positional_to_numpy"):
-      array_a_as_np = np.array(array_a)
-      self.assertEqual(array_a_as_np.shape, (6,))
-
-    with self.subTest("convert_positional_to_jax"):
-      array_a_as_jax = jnp.array(array_a)
-      self.assertEqual(array_a_as_jax.shape, (6,))
-
-    with self.subTest("conversion_failure_with_names"):
-      with self.assertRaisesRegex(
-          ValueError,
-          re.escape(
-              "Only NamedArray(View)s with no named axes can be converted to"
-              " numpy arrays"
-          ),
-      ):
-        _ = np.array(array_a.tag("foo"))
-
-    with self.subTest("conversion_failure_with_names_jax"):
-      with self.assertRaisesRegex(
-          ValueError,
-          re.escape(
-              "Only NamedArray(View)s with no named axes can be converted to"
-              " JAX arrays"
-          ),
-      ):
-        _ = jnp.array(array_a.tag("foo"))
-
   def test_convenience_constructors(self):
     with self.subTest("zeros"):
       chex.assert_trees_all_equal(