feat: Update JAX backend array_type and array_subtype to jax.Array (#2079)

* Update JAX backend array_type and array_subtype to jax.Array, which is a unified
  array type introduced in jax v0.4.1 that subsumes DeviceArray, ShardedDeviceArray,
  and GlobalDeviceArray types in JAX.
   - c.f. https://github.com/google/jax/releases/tag/jax-v0.4.1
* Update the lower bound of jax and jaxlib to v0.4.1 to ensure support for jax.Array.
* Update lower bound of the supported scipy versions to v1.5.0 as required by jax v0.4.1.
* Update the JAX backend docstring examples to use Array as the array type.

Author: matthewfeickert

setup.cfg

@@ -34,7 +34,7 @@ package_dir =
 include_package_data = True
 python_requires = >=3.8
 install_requires =
-    scipy>=1.3.2 # requires numpy, which is required by pyhf and tensorflow
+    scipy>=1.5.0 # requires numpy, which is required by pyhf and tensorflow
     click>=8.0.0  # for console scripts
     tqdm>=4.56.0  # for readxml
     jsonschema>=4.15.0  # for utils

setup.py

@@ -7,7 +7,7 @@
         'tensorflow-probability>=0.11.0',  # c.f. PR #1657
     ],
     'torch': ['torch>=1.10.0'],  # c.f. PR #1657
-    'jax': ['jax>=0.2.10', 'jaxlib>=0.1.61,!=0.1.68'],  # c.f. PR #1962, Issue #1501
+    'jax': ['jax>=0.4.1', 'jaxlib>=0.4.1'],  # c.f. PR #2079
     'xmlio': ['uproot>=4.1.1'],  # c.f. PR #1567
     'minuit': ['iminuit>=2.7.0'],  # c.f. PR #1895
 }

src/pyhf/tensor/jax_backend.py

@@ -2,6 +2,7 @@
 
 config.update('jax_enable_x64', True)
 
+from jax import Array
 import jax.numpy as jnp
 from jax.scipy.special import gammaln, xlogy
 from jax.scipy import special
@@ -54,10 +55,10 @@ class jax_backend:
     __slots__ = ['name', 'precision', 'dtypemap', 'default_do_grad']
 
     #: The array type for jax
-    array_type = jnp.DeviceArray
+    array_type = Array
 
     #: The array content type for jax
-    array_subtype = jnp.DeviceArray
+    array_subtype = Array
 
     def __init__(self, **kwargs):
         self.name = 'jax'
@@ -84,7 +85,7 @@ def clip(self, tensor_in, min_value, max_value):
             >>> pyhf.set_backend("jax")
             >>> a = pyhf.tensorlib.astensor([-2, -1, 0, 1, 2])
             >>> pyhf.tensorlib.clip(a, -1, 1)
-            DeviceArray([-1., -1.,  0.,  1.,  1.], dtype=float64)
+            Array([-1., -1.,  0.,  1.,  1.], dtype=float64)
 
         Args:
             tensor_in (:obj:`tensor`): The input tensor object
@@ -106,8 +107,7 @@ def erf(self, tensor_in):
             >>> pyhf.set_backend("jax")
             >>> a = pyhf.tensorlib.astensor([-2., -1., 0., 1., 2.])
             >>> pyhf.tensorlib.erf(a)
-            DeviceArray([-0.99532227, -0.84270079,  0.        ,  0.84270079,
-                          0.99532227], dtype=float64)
+            Array([-0.99532227, -0.84270079,  0.        ,  0.84270079,  0.99532227],      dtype=float64)
 
         Args:
             tensor_in (:obj:`tensor`): The input tensor object
@@ -127,7 +127,7 @@ def erfinv(self, tensor_in):
             >>> pyhf.set_backend("jax")
             >>> a = pyhf.tensorlib.astensor([-2., -1., 0., 1., 2.])
             >>> pyhf.tensorlib.erfinv(pyhf.tensorlib.erf(a))
-            DeviceArray([-2., -1.,  0.,  1.,  2.], dtype=float64)
+            Array([-2., -1.,  0.,  1.,  2.], dtype=float64)
 
         Args:
             tensor_in (:obj:`tensor`): The input tensor object
@@ -147,8 +147,8 @@ def tile(self, tensor_in, repeats):
             >>> pyhf.set_backend("jax")
             >>> a = pyhf.tensorlib.astensor([[1.0], [2.0]])
             >>> pyhf.tensorlib.tile(a, (1, 2))
-            DeviceArray([[1., 1.],
-                         [2., 2.]], dtype=float64)
+            Array([[1., 1.],
+                   [2., 2.]], dtype=float64)
 
         Args:
             tensor_in (:obj:`tensor`): The tensor to be repeated
@@ -171,7 +171,7 @@ def conditional(self, predicate, true_callable, false_callable):
             >>> a = tensorlib.astensor([4])
             >>> b = tensorlib.astensor([5])
             >>> tensorlib.conditional((a < b)[0], lambda: a + b, lambda: a - b)
-            DeviceArray([9.], dtype=float64)
+            Array([9.], dtype=float64)
 
         Args:
             predicate (:obj:`scalar`): The logical condition that determines which callable to evaluate
@@ -213,16 +213,16 @@ def astensor(self, tensor_in, dtype="float"):
             >>> pyhf.set_backend("jax")
             >>> tensor = pyhf.tensorlib.astensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
             >>> tensor
-            DeviceArray([[1., 2., 3.],
-                         [4., 5., 6.]], dtype=float64)
+            Array([[1., 2., 3.],
+                   [4., 5., 6.]], dtype=float64)
             >>> type(tensor) # doctest:+ELLIPSIS
-            <class '...DeviceArray'>
+            <class '...Array'>
 
         Args:
             tensor_in (Number or Tensor): Tensor object
 
         Returns:
-            `jaxlib.xla_extension.DeviceArray`: A multi-dimensional, fixed-size homogeneous array.
+            `jaxlib.xla_extension.Array`: A multi-dimensional, fixed-size homogeneous array.
         """
         # TODO: Remove doctest:+ELLIPSIS when JAX API stabilized
         try:
@@ -294,9 +294,9 @@ def percentile(self, tensor_in, q, axis=None, interpolation="linear"):
             >>> pyhf.set_backend("jax")
             >>> a = pyhf.tensorlib.astensor([[10, 7, 4], [3, 2, 1]])
             >>> pyhf.tensorlib.percentile(a, 50)
-            DeviceArray(3.5, dtype=float64)
+            Array(3.5, dtype=float64)
             >>> pyhf.tensorlib.percentile(a, 50, axis=1)
-            DeviceArray([7., 2.], dtype=float64)
+            Array([7., 2.], dtype=float64)
 
         Args:
             tensor_in (`tensor`): The tensor containing the data
@@ -355,7 +355,7 @@ def simple_broadcast(self, *args):
             ...   pyhf.tensorlib.astensor([1]),
             ...   pyhf.tensorlib.astensor([2, 3, 4]),
             ...   pyhf.tensorlib.astensor([5, 6, 7]))
-            [DeviceArray([1., 1., 1.], dtype=float64), DeviceArray([2., 3., 4.], dtype=float64), DeviceArray([5., 6., 7.], dtype=float64)]
+            [Array([1., 1., 1.], dtype=float64), Array([2., 3., 4.], dtype=float64), Array([5., 6., 7.], dtype=float64)]
 
         Args:
             args (Array of Tensors): Sequence of arrays
@@ -381,13 +381,13 @@ def ravel(self, tensor):
             >>> pyhf.set_backend("jax")
             >>> tensor = pyhf.tensorlib.astensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
             >>> pyhf.tensorlib.ravel(tensor)
-            DeviceArray([1., 2., 3., 4., 5., 6.], dtype=float64)
+            Array([1., 2., 3., 4., 5., 6.], dtype=float64)
 
         Args:
             tensor (Tensor): Tensor object
 
         Returns:
-            `jaxlib.xla_extension.DeviceArray`: A flattened array.
+            `jaxlib.xla_extension.Array`: A flattened array.
         """
         return jnp.ravel(tensor)
 
@@ -441,11 +441,11 @@ def poisson(self, n, lam):
             >>> import pyhf
             >>> pyhf.set_backend("jax")
             >>> pyhf.tensorlib.poisson(5., 6.)
-            DeviceArray(0.16062314, dtype=float64, weak_type=True)
+            Array(0.16062314, dtype=float64, weak_type=True)
             >>> values = pyhf.tensorlib.astensor([5., 9.])
             >>> rates = pyhf.tensorlib.astensor([6., 8.])
             >>> pyhf.tensorlib.poisson(values, rates)
-            DeviceArray([0.16062314, 0.12407692], dtype=float64)
+            Array([0.16062314, 0.12407692], dtype=float64)
 
         Args:
             n (:obj:`tensor` or :obj:`float`): The value at which to evaluate the approximation to the Poisson distribution p.m.f.
@@ -484,12 +484,12 @@ def normal(self, x, mu, sigma):
             >>> import pyhf
             >>> pyhf.set_backend("jax")
             >>> pyhf.tensorlib.normal(0.5, 0., 1.)
-            DeviceArray(0.35206533, dtype=float64, weak_type=True)
+            Array(0.35206533, dtype=float64, weak_type=True)
             >>> values = pyhf.tensorlib.astensor([0.5, 2.0])
             >>> means = pyhf.tensorlib.astensor([0., 2.3])
             >>> sigmas = pyhf.tensorlib.astensor([1., 0.8])
             >>> pyhf.tensorlib.normal(values, means, sigmas)
-            DeviceArray([0.35206533, 0.46481887], dtype=float64)
+            Array([0.35206533, 0.46481887], dtype=float64)
 
         Args:
             x (:obj:`tensor` or :obj:`float`): The value at which to evaluate the Normal distribution p.d.f.
@@ -510,10 +510,10 @@ def normal_cdf(self, x, mu=0, sigma=1):
             >>> import pyhf
             >>> pyhf.set_backend("jax")
             >>> pyhf.tensorlib.normal_cdf(0.8)
-            DeviceArray(0.7881446, dtype=float64)
+            Array(0.7881446, dtype=float64)
             >>> values = pyhf.tensorlib.astensor([0.8, 2.0])
             >>> pyhf.tensorlib.normal_cdf(values)
-            DeviceArray([0.7881446 , 0.97724987], dtype=float64)
+            Array([0.7881446 , 0.97724987], dtype=float64)
 
         Args:
             x (:obj:`tensor` or :obj:`float`): The observed value of the random variable to evaluate the CDF for
@@ -536,7 +536,7 @@ def poisson_dist(self, rate):
             >>> values = pyhf.tensorlib.astensor([4, 9])
             >>> poissons = pyhf.tensorlib.poisson_dist(rates)
             >>> poissons.log_prob(values)
-            DeviceArray([-1.74030218, -2.0868536 ], dtype=float64)
+            Array([-1.74030218, -2.0868536 ], dtype=float64)
 
         Args:
             rate (:obj:`tensor` or :obj:`float`): The mean of the Poisson distribution (the expected number of events)
@@ -558,7 +558,7 @@ def normal_dist(self, mu, sigma):
             >>> values = pyhf.tensorlib.astensor([4, 9])
             >>> normals = pyhf.tensorlib.normal_dist(means, stds)
             >>> normals.log_prob(values)
-            DeviceArray([-1.41893853, -2.22579135], dtype=float64)
+            Array([-1.41893853, -2.22579135], dtype=float64)
 
         Args:
             mu (:obj:`tensor` or :obj:`float`): The mean of the Normal distribution
@@ -579,8 +579,8 @@ def to_numpy(self, tensor_in):
             >>> pyhf.set_backend("jax")
             >>> tensor = pyhf.tensorlib.astensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
             >>> tensor
-            DeviceArray([[1., 2., 3.],
-                         [4., 5., 6.]], dtype=float64)
+            Array([[1., 2., 3.],
+                   [4., 5., 6.]], dtype=float64)
             >>> numpy_ndarray = pyhf.tensorlib.to_numpy(tensor)
             >>> numpy_ndarray
             array([[1., 2., 3.],
@@ -606,12 +606,12 @@ def transpose(self, tensor_in):
             >>> pyhf.set_backend("jax")
             >>> tensor = pyhf.tensorlib.astensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
             >>> tensor
-            DeviceArray([[1., 2., 3.],
-                         [4., 5., 6.]], dtype=float64)
+            Array([[1., 2., 3.],
+                   [4., 5., 6.]], dtype=float64)
             >>> pyhf.tensorlib.transpose(tensor)
-            DeviceArray([[1., 4.],
-                         [2., 5.],
-                         [3., 6.]], dtype=float64)
+            Array([[1., 4.],
+                   [2., 5.],
+                   [3., 6.]], dtype=float64)
 
         Args:
             tensor_in (:obj:`tensor`): The input tensor object.

tests/constraints.txt

@@ -1,5 +1,5 @@
 # core
-scipy==1.3.2  # c.f. PR #2044
+scipy==1.5.0  # c.f. PR #2079
 click==8.0.0  # c.f. PR #1958, #1909
 tqdm==4.56.0
 jsonschema==4.15.0  # c.f. PR #1979
@@ -19,5 +19,5 @@ torch==1.10.0
 # Use Google Cloud Storage buckets for long term wheel support
 # c.f. https://github.com/google/jax/discussions/7608#discussioncomment-1269342
 --find-links https://storage.googleapis.com/jax-releases/jax_releases.html
-jax==0.2.10
-jaxlib==0.1.61  # c.f. PR #1962
+jax==0.4.1  # c.f. PR #2079
+jaxlib==0.4.1  # c.f. PR #2079