Remove the deprecated import jax.tree_map (pyro-ppl#1775)

* remove deprecated import jax.tree_map

* fix enable_x64 logic

numpyro/contrib/einstein/mixture_guide_predictive.py

@@ -5,8 +5,8 @@
 from functools import partial
 from typing import Optional
 
-from jax import numpy as jnp, random, tree_map, vmap
-from jax.tree_util import tree_flatten
+from jax import numpy as jnp, random, vmap
+from jax.tree_util import tree_flatten, tree_map
 
 from numpyro.handlers import substitute
 from numpyro.infer import Predictive

numpyro/distributions/batch_util.py

@@ -5,8 +5,8 @@
 from functools import singledispatch
 from typing import Union
 
-from jax import tree_map
 import jax.numpy as jnp
+from jax.tree_util import tree_map
 
 from numpyro.distributions import constraints
 from numpyro.distributions.conjugate import (

numpyro/util.py

@@ -45,7 +45,7 @@ def enable_x64(use_x64=True):
     """
     if not use_x64:
         use_x64 = os.getenv("JAX_ENABLE_X64", 0)
-    jax.config.update("jax_enable_x64", use_x64)
+    jax.config.update("jax_enable_x64", bool(use_x64))
 
 
 def set_platform(platform=None):

test/test_constraints.py

@@ -5,8 +5,9 @@
 
 import pytest
 
-from jax import jit, tree_map, vmap
+from jax import jit, vmap
 import jax.numpy as jnp
+from jax.tree_util import tree_map
 
 from numpyro.distributions import constraints
 

test/test_transforms.py

@@ -8,8 +8,9 @@
 import numpy as np
 import pytest
 
-from jax import jacfwd, jit, random, tree_map, vmap
+from jax import jacfwd, jit, random, vmap
 import jax.numpy as jnp
+from jax.tree_util import tree_map
 
 from numpyro.distributions import constraints
 from numpyro.distributions.flows import (