Repository: google/jax
Issue: BUG: jax.Array allows attribute setting

The new `jax.Array` implementation allows arbitrary attributes to be added to array objects. This was not the case for `DeviceArray`:
```python
# test.py
import jax.numpy as jnp
x = jnp.arange(4)
x.custom_attr = 5
print(x.custom_attr)
```
```
$ JAX_ARRAY=1 python test.py
5

$ JAX_ARRAY=0 python test.py 
Traceback (most recent call last):
  File "tmp.py", line 3, in <module>
    x.custom_attr = 5
AttributeError: 'DeviceArray' object has no attribute 'custom_attr'
```
Why is this a problem? Users who depend on this new behavior will end up with JAX code that doesn't transform correctly. For example:
```python
def f(x):
  x.custom_attr = 4
  return x

print(f(jnp.arange(4)).custom_attr)
# 4

print(jit(f)(jnp.arange(4)).custom_attr)
# AttributeError: 'DynamicJaxprTracer' object has no attribute 'custom_attr'
```
We should discourage this kind of bug by disallowing such attributes entirely, as in the previous `DeviceArray` implementation.

Changed files:
tests/array_test.py, jax/_src/array.py

diff --git a/jax/_src/array.py b/jax/_src/array.py
index 2542d3d8b..4a278b5b2 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -16,7 +16,6 @@ from __future__ import annotations
 
 import operator as op
 import numpy as np
-import functools
 from typing import Sequence, Tuple, Callable, Union, Optional, cast, List
 
 from jax._src import abstract_arrays
@@ -109,7 +108,6 @@ def _single_device_array_from_buf(buf, committed):
 
 @use_cpp_class(xc.ArrayImpl)
 class ArrayImpl(basearray.Array):
-  # TODO(yashkatariya): Add __slots__ here.
 
   aval: core.ShapedArray
   _sharding: Sharding
@@ -122,7 +120,9 @@ class ArrayImpl(basearray.Array):
   def __init__(self, aval: core.ShapedArray, sharding: Sharding,
                arrays: Union[Sequence[DeviceArray], Sequence[ArrayImpl]],
                committed: bool, _skip_checks: bool = False):
+    ##########################################################################
     # NOTE: the actual implementation of the constructor is moved to C++.
+    ##########################################################################
 
     self.aval = aval
     self._sharding = sharding
@@ -135,6 +135,7 @@ class ArrayImpl(basearray.Array):
     # for what committed means.
     self._committed = committed
     self._npy_value = None
+    self._cached_addressable_shards = None
 
     # Don't rearrange if skip_checks is enabled because this assumes that the
     # input buffers are already arranged properly. This usually happens when
@@ -341,7 +342,7 @@ class ArrayImpl(basearray.Array):
     else:
       return f"{prefix}{self.shape}, {dtype_str}"
 
-  @functools.cached_property
+  @property
   def is_fully_addressable(self) -> bool:
     return self.sharding.is_fully_addressable
 
@@ -412,16 +413,19 @@ class ArrayImpl(basearray.Array):
     self._check_if_deleted()
     return _single_device_array_from_buf(self._arrays[index], self._committed)
 
-  @functools.cached_property
-  def addressable_shards(self) -> Sequence[Shard]:
+  @property
+  def addressable_shards(self):
     self._check_if_deleted()
-    out = []
-    for db in self._arrays:
-      # Wrap the device arrays in `Array` until C++ returns an Array instead
-      # of a DA.
-      array = _single_device_array_from_buf(db, self._committed)
-      out.append(Shard(db.device(), self.sharding, self.shape, array))
-    return out
+
+    if getattr(self, '_cached_addressable_shards', None) is None:
+      out = []
+      for db in self._arrays:
+        # Wrap the device arrays in `Array` until C++ returns an Array instead
+        # of a DA.
+        array = _single_device_array_from_buf(db, self._committed)
+        out.append(Shard(db.device(), self.sharding, self.shape, array))
+      self._cached_addressable_shards = out
+    return self._cached_addressable_shards
 
   @property
   def global_shards(self) -> Sequence[Shard]:
diff --git a/tests/array_test.py b/tests/array_test.py
index 114c74edc..d78e56dce 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -26,6 +26,7 @@ from jax._src import dispatch
 from jax._src import test_util as jtu
 from jax._src.lib import xla_client as xc
 from jax._src.lib import xla_bridge as xb
+from jax._src.lib import xla_extension_version
 from jax._src.util import prod, safe_zip
 from jax.interpreters import pxla
 from jax.experimental.pjit import pjit
@@ -643,6 +644,14 @@ class JaxArrayTest(jtu.JaxTestCase):
     self.assertEqual(shard_size * len(a.global_shards),
                      a.on_device_size_in_bytes())
 
+  def test_disallow_setting_attr(self):
+    if xla_extension_version < 116:
+      self.skipTest("This test requires xla_extension_version >= 116")
+    x = jnp.arange(4)
+    with self.assertRaises(AttributeError):
+      x.custom_attr = 5
+
+
 @jtu.with_config(jax_array=True)
 class ShardingTest(jtu.JaxTestCase):