Fix package exports

jax/__init__.py

@@ -106,11 +106,11 @@
   tree_structure,
   tree_transpose,
   tree_unflatten,
-  value_and_grad,
-  vjp,
-  vmap,
+  value_and_grad as value_and_grad,
+  vjp as vjp,
+  vmap as vmap,
   xla,  # TODO(phawkins): update users to avoid this.
-  xla_computation,
+  xla_computation as xla_computation,
 )
 from functools import partial as partial  # TODO(phawkins): remove this export.
 from .experimental.maps import soft_pmap as soft_pmap

jax/experimental/__init__.py

@@ -13,6 +13,12 @@
 # limitations under the License.
 
 # flake8: noqa: F401
-from ..interpreters.sharded_jit import (sharded_jit, PartitionSpec,
-                                        with_sharding_constraint)
-from .x64_context import enable_x64, disable_x64
+from ..interpreters.sharded_jit import (
+  sharded_jit as sharded_jit,
+  PartitionSpec as PartitionSpec,
+  with_sharding_constraint as with_sharding_constraint,
+)
+from .x64_context import (
+  enable_x64 as enable_x64,
+  disable_x64 as disable_x64,
+)

jax/experimental/sparse/__init__.py

@@ -183,44 +183,47 @@
 """
 
 # flake8: noqa: F401
-from .ad import grad, value_and_grad
+from .ad import (
+    grad as grad,
+    value_and_grad as value_and_grad,
+)
 from .ops import (
-    bcoo_dot_general,
-    bcoo_dot_general_p,
-    bcoo_dot_general_sampled,
-    bcoo_dot_general_sampled_p,
-    bcoo_extract,
-    bcoo_extract_p,
-    bcoo_fromdense,
-    bcoo_fromdense_p,
-    bcoo_reduce_sum,
-    bcoo_rdot_general,
-    bcoo_spdot_general,
-    bcoo_spdot_general_p,
-    bcoo_todense,
-    bcoo_todense_p,
-    bcoo_transpose,
-    bcoo_transpose_p,
-    coo_fromdense,
-    coo_fromdense_p,
-    coo_matmat,
-    coo_matmat_p,
-    coo_matvec,
-    coo_matvec_p,
-    coo_todense,
-    coo_todense_p,
-    csr_fromdense,
-    csr_fromdense_p,
-    csr_matmat,
-    csr_matmat_p,
-    csr_matvec,
-    csr_matvec_p,
-    csr_todense,
-    csr_todense_p,
-    COO,
-    CSC,
-    CSR,
-    BCOO,
+    bcoo_dot_general as bcoo_dot_general,
+    bcoo_dot_general_p as bcoo_dot_general_p,
+    bcoo_dot_general_sampled as bcoo_dot_general_sampled,
+    bcoo_dot_general_sampled_p as bcoo_dot_general_sampled_p,
+    bcoo_extract as bcoo_extract,
+    bcoo_extract_p as bcoo_extract_p,
+    bcoo_fromdense as bcoo_fromdense,
+    bcoo_fromdense_p as bcoo_fromdense_p,
+    bcoo_reduce_sum as bcoo_reduce_sum,
+    bcoo_rdot_general as bcoo_rdot_general,
+    bcoo_spdot_general as bcoo_spdot_general,
+    bcoo_spdot_general_p as bcoo_spdot_general_p,
+    bcoo_todense as bcoo_todense,
+    bcoo_todense_p as bcoo_todense_p,
+    bcoo_transpose as bcoo_transpose,
+    bcoo_transpose_p as bcoo_transpose_p,
+    coo_fromdense as coo_fromdense,
+    coo_fromdense_p as coo_fromdense_p,
+    coo_matmat as coo_matmat,
+    coo_matmat_p as coo_matmat_p,
+    coo_matvec as coo_matvec,
+    coo_matvec_p as coo_matvec_p,
+    coo_todense as coo_todense,
+    coo_todense_p as coo_todense_p,
+    csr_fromdense as csr_fromdense,
+    csr_fromdense_p as csr_fromdense_p,
+    csr_matmat as csr_matmat,
+    csr_matmat_p as csr_matmat_p,
+    csr_matvec as csr_matvec,
+    csr_matvec_p as csr_matvec_p,
+    csr_todense as csr_todense,
+    csr_todense_p as csr_todense_p,
+    COO as COO,
+    CSC as CSC,
+    CSR as CSR,
+    BCOO as BCOO,
 )
 
-from .transform import sparsify
+from .transform import sparsify as sparsify

jax/nn/__init__.py

@@ -16,8 +16,8 @@
 
 # flake8: noqa: F401
 
-from jax.numpy import tanh
-from . import initializers
+from jax.numpy import tanh as tanh
+from . import initializers as initializers
 from jax._src.nn.functions import (
   celu as celu,
   elu as elu,

jax/scipy/stats/__init__.py

@@ -13,21 +13,21 @@
 # limitations under the License.
 
 # flake8: noqa: F401
-from . import bernoulli
-from . import beta
-from . import cauchy
-from . import dirichlet
-from . import expon
-from . import gamma
-from . import geom
-from . import laplace
-from . import logistic
-from . import multivariate_normal
-from . import nbinom
-from . import norm
-from . import pareto
-from . import poisson
-from . import t
-from . import uniform
-from . import chi2
-from . import betabinom
+from . import bernoulli as bernoulli
+from . import beta as beta
+from . import cauchy as cauchy
+from . import dirichlet as dirichlet
+from . import expon as expon
+from . import gamma as gamma
+from . import geom as geom
+from . import laplace as laplace
+from . import logistic as logistic
+from . import multivariate_normal as multivariate_normal
+from . import nbinom as nbinom
+from . import norm as norm
+from . import pareto as pareto
+from . import poisson as poisson
+from . import t as t
+from . import uniform as uniform
+from . import chi2 as chi2
+from . import betabinom as betabinom