Fix import bug issue with jax.stages.ArgInfo in past JAX versions.

.github/workflows/tests.yaml

@@ -24,10 +24,10 @@ jobs:
  access_token: ${{ github.token }}
  if: ${{github.ref != 'refs/head/main'}}
  - uses: actions/checkout@v3
- - name: Set up Python 3.8
+ - name: Set up Python 3.10
  uses: actions/setup-python@v4
  with:
- python-version: 3.8
+ python-version: 3.10
  - uses: pre-commit/action@v3.0.0
 
  unit_tests:

jax_scalify/core/datatype.py

@@ -204,7 +204,7 @@ def as_scaled_array_base(
  return ScaledArray(val / scale.astype(val.dtype), scale) # type:ignore
 
  # TODO: fix bug when scale is not 1.
- raise NotImplementedError(f"Constructing `ScaledArray` from {val} and {scale} not supported.")
+ raise NotImplementedError(f"Constructing `ScaledArray` from {val} and {scale} not supported.") # type:ignore
  # return scaled_array_base(val, scale)
 
 

jax_scalify/core/typing.py

@@ -1,5 +1,5 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
-from typing import Any
+from typing import Any, Tuple
 
 # import chex
 import jax
@@ -10,12 +10,21 @@
 try:
  from jax import Array
 
- ArrayTypes = (Array, jax.stages.ArgInfo)
+ ArrayTypes: Tuple[Any, ...] = (Array,)
 except ImportError:
  from jaxlib.xla_extension import DeviceArray as Array
 
+ # Older version of JAX <0.4
  ArrayTypes = (Array, jax.interpreters.partial_eval.DynamicJaxprTracer)
 
+try:
+ from jax.stages import ArgInfo
+
+ # Additional ArgInfo in recent JAX versions.
+ ArrayTypes = (*ArrayTypes, ArgInfo)
+except ImportError:
+ pass
+
 
 def get_numpy_api(val: Any) -> Any:
  """Get the Numpy API corresponding to an array.