Closed matthewfeickert closed 1 year ago
This particular issue can be fixed with just
$ git diff
diff --git a/src/pyhf/tensor/jax_backend.py b/src/pyhf/tensor/jax_backend.py
index 2e05039f..9cd46a58 100644
--- a/src/pyhf/tensor/jax_backend.py
+++ b/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'
=
but depending on how much else needs to get updated I'm not sure if this will be just a patch fix or if this should be viewed as part of the work for the next minor release given that the only thing that fails is schema validation checks, that would pass on an older release.
edit: At runtime this is basically fine, so I think this can be moved to the next minor release.
Honestly, this is going to impact one person for right now @phinate but that's probably ok.
Honestly, this is going to impact one person for right now @phinate but that's probably ok.
goes to cry in corner :'(
But in all seriousness, this is unlikely to affect me as I'm still bypassing validation right now in experiments (something about my fork seemed to necessitate it), and also I'm not actively working on anything either. So go for it :)
Thanks for the feedback @phinate. :+1: I'll go ahead and merge in PR #2079 then, but after your thesis defense if you need any new releases of pyhf
to help make things easier on the neos
or relaxed
side just let us know and we can work with you on that. I want to make it easier for us to cut releases as needed.
In JAX v0.4.1
This causes
test_schema_tensor_type_allowed
to failAs this is failing on schema tests this might require a patch release.