Move jax.api to jax._src.api.

PiperOrigin-RevId: 368233837
google · Apr 13, 2021 · 26e9ebc · 26e9ebc
1 parent be8cba3
commit 26e9ebc
Show file tree

Hide file tree

Showing 39 changed files with 2,620 additions and 2,560 deletions.
diff --git a/docs/notebooks/How_JAX_primitives_work.ipynb b/docs/notebooks/How_JAX_primitives_work.ipynb
@@ -85,7 +85,7 @@
    ],
    "source": [
     "from jax import lax\n",
-    "from jax import api\n",
+    "from jax._src import api\n",
     "\n",
     "def multiply_add_lax(x, y, z):\n",
     "  \"\"\"Implementation of multiply-add using the jax.lax primitives.\"\"\"\n",

diff --git a/docs/notebooks/How_JAX_primitives_work.md b/docs/notebooks/How_JAX_primitives_work.md
@@ -64,7 +64,7 @@ defined in the `jax.lax` module:
 :outputId: 3fb1c8a7-7a4c-4a3a-f7ff-37b7dc740528
 
 from jax import lax
-from jax import api
+from jax._src import api
 
 def multiply_add_lax(x, y, z):
   """Implementation of multiply-add using the jax.lax primitives."""

diff --git a/examples/advi.py b/examples/advi.py
@@ -22,7 +22,7 @@
 from functools import partial
 import matplotlib.pyplot as plt
 
-from jax.api import jit, grad, vmap
+from jax import jit, grad, vmap
 from jax import random
 from jax.experimental import optimizers
 import jax.numpy as jnp

diff --git a/jax/__init__.py b/jax/__init__.py
@@ -33,7 +33,7 @@
 from .config import (config, enable_checks, check_tracer_leaks, checking_leaks,
                      debug_nans, debug_infs, log_compiles,
                      default_matmul_precision, numpy_rank_promotion)
-from .api import (
+from ._src.api import (
   ad,  # TODO(phawkins): update users to avoid this.
   checkpoint,
   closure_convert,
@@ -99,6 +99,7 @@
 
 # These submodules are separate because they are in an import cycle with
 # jax and rely on the names imported above.
+from . import api
 from . import dtypes
 from . import errors
 from . import image