TMP

jakevdp · jakevdp · commit 1e5ec7815e77 · 2022-09-19T15:55:15.000-07:00
diff --git a/jax/core.py b/jax/core.py
@@ -50,6 +50,7 @@
 from jax._src.lib import jax_jit
 from jax._src import traceback_util
 from jax._src.typing import DimSize, Shape
+from jax._src import typing
 traceback_util.register_exclusion(__file__)
 
 zip, unsafe_zip = safe_zip, zip
@@ -530,9 +531,9 @@ def escaped_tracer_error(tracer, detail=None):
     msg += f'Detail: {detail}'
   return UnexpectedTracerError(msg)
 
-class Tracer:
+class Tracer(typing.Array):
   __array_priority__ = 1000
-  __slots__ = ['_trace', '__weakref__', '_line_info']
+  __slots__ = ['_trace', '_line_info']
 
   def __array__(self, *args, **kw):
     raise TracerArrayConversionError(self)
diff --git a/jax/experimental/array.py b/jax/experimental/array.py
@@ -95,7 +95,7 @@ def _reconstruct_array(fun, args, arr_state, aval_state):
 
 
 @pxla.use_cpp_class(xc.Array if xc._version >= 92 else None)
-class Array:
+class Array(typing.Array):
   # TODO(yashkatariya): Add __slots__ here.
 
   @pxla.use_cpp_method
diff --git a/tests/typing_test.py b/tests/typing_test.py
@@ -17,14 +17,17 @@
 This test is meant to be both a runtime test and a static type annotation test,
 so it should be checked with pytype/mypy as well as being run with pytest.
 """
-from typing import Union
+from typing import Any, Optional, Union
 
 import jax
+from jax import core
 from jax._src import test_util as jtu
 from jax._src import typing
 from jax import lax
 import jax.numpy as jnp
 
+from jax.experimental.array import Array as ArrayImpl
+
 from absl.testing import absltest
 import numpy as np
 
@@ -109,6 +112,21 @@ def is_array(x: typing.ArrayLike) -> Union[bool, typing.Array]:
     self.assertTrue(jax.jit(is_array)(x))
     self.assertTrue(jnp.all(jax.vmap(is_array)(x)))
 
+  def testAnnotations(self):
+    # This test is mainly meant for static type checking: we want to ensure that
+    # Tracer and ArrayImpl are valid as array.Array.
+
+    def f(x: Any) -> Optional[typing.Array]:
+      if isinstance(x, core.Tracer):
+        return x
+      elif isinstance(x, ArrayImpl):
+        return x
+      else:
+        return None
+
+    x = jnp.arange(10)
+    y = f(x)
+    self.assertArraysEqual(x, y)
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
