fix(backport): Access jax.config from jax (#2423)

matthewfeickert · web-flow · commit 9c9fc876a009 · 2024-01-05T01:46:17.000-06:00
* Backport PR #2376 * As of jax and jaxlib v0.4.20 accessing jax.config from the jax.config submodule is deprecated and it should be accessed from the jax top level API instead. ``` >>> from jax.config import config <stdin>:1: DeprecationWarning: Accessing jax.config via the jax.config submodule is deprecated ```
diff --git a/src/pyhf/tensor/jax_backend.py b/src/pyhf/tensor/jax_backend.py
@@ -1,4 +1,4 @@
-from jax.config import config
+from jax import config
 
 config.update('jax_enable_x64', True)
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-from jax.config import config`
	`1`	`+from jax import config`
`2`	`2`
`3`	`3`	`config.update('jax_enable_x64', True)`
`4`	`4`