Export jax.tree* to avoid breaking users.

jax-ml · May 7, 2020 · 68daf38 · 68daf38
1 parent e892ad0
commit 68daf38
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/docs/notebooks/How_JAX_primitives_work.ipynb b/docs/notebooks/How_JAX_primitives_work.ipynb
@@ -700,7 +700,7 @@
         "\n",
         "# Now we register the XLA compilation rule with JAX\n",
         "# TODO: for GPU? and TPU?\n",
-        "from jax import xla\n",
+        "from jax.interpreters import xla\n",
         "xla.backend_specific_translations['cpu'][multiply_add_p] = multiply_add_xla_translation"
       ],
       "execution_count": 0,

diff --git a/jax/__init__.py b/jax/__init__.py
@@ -50,6 +50,16 @@
   ShapedArray,
   ShapeDtypeStruct,
   soft_pmap,
+  # TODO(phawkins): hide tree* functions from jax, update callers to use
+  # jax.tree_util.
+  treedef_is_leaf,
+  tree_flatten,
+  tree_leaves,
+  tree_map,
+  tree_multimap,
+  tree_structure
+  tree_transpose,
+  tree_unflatten,
   value_and_grad,
   vjp,
   vmap,
@@ -62,7 +72,6 @@
 def _init():
   import os
   os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '1')
-  del os
 
   import jax.numpy  # side-effecting import sets up operator overloads