Document why 'import name as name' is used

jax/__init__.py

@@ -35,6 +35,9 @@ del _cloud_tpu_init
 from jax import config as _config_module
 del _config_module
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.basearray import Array as Array
 
 from jax._src.config import (

jax/custom_derivatives.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.custom_derivatives import (
   _initial_style_jaxpr,
   _sum_tangents,

jax/dtypes.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.dtypes import (
     _jax_types,  # TODO(phawkins): fix users and remove?
     bfloat16 as bfloat16,

jax/errors.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.errors import (
   JAXTypeError as JAXTypeError,
   JAXIndexError as JAXIndexError,

jax/experimental/__init__.py

@@ -14,6 +14,10 @@
 
 # TODO(https://github.com/google/jax/issues/13487): Remove PartitionSpec in
 # 3 months from `jax.experimental.PartitionSpec`.
+
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax.interpreters.pxla import PartitionSpec as PartitionSpec
 from jax.experimental.x64_context import (
   enable_x64 as enable_x64,

jax/experimental/checkify.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.checkify import (
     Error as Error,
     ErrorCategory as ErrorCategory,

jax/experimental/sparse/__init__.py

@@ -183,6 +183,9 @@ To fit the same model on sparse data, we can apply the :func:`sparsify` transfor
      -0.670236    0.03132951 -0.05356663]
 """
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax.experimental.sparse.ad import (
     grad as grad,
     value_and_grad as value_and_grad,

jax/image/__init__.py

@@ -20,6 +20,9 @@ JAX, such as `PIX`_.
 .. _PIX: https://github.com/deepmind/dm_pix
 """
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.image.scale import (
   resize as resize,
   ResizeMethod as ResizeMethod,

jax/lax/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.lax.lax import (
   DotDimensionNumbers as DotDimensionNumbers,
   Precision as Precision,

jax/nn/__init__.py

@@ -14,6 +14,9 @@
 
 """Common functions for neural network libraries."""
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax.numpy import tanh as tanh
 from jax.nn import initializers as initializers
 from jax._src.nn.functions import (

jax/nn/initializers.py

@@ -17,6 +17,9 @@ Common neural network layer initializers, consistent with definitions
 used in Keras and Sonnet.
 """
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.nn.initializers import (
   constant as constant,
   Initializer as Initializer,

jax/numpy/fft.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.numpy.fft import (
   ifft as ifft,
   ifft2 as ifft2,

jax/numpy/linalg.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.numpy.linalg import (
   cholesky as cholesky,
   det as det,

jax/ops/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.ops.scatter import (
   segment_sum as segment_sum,
   segment_prod as segment_prod,

jax/prng.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.prng import (
   PRNGImpl as PRNGImpl,
   seed_with_impl as seed_with_impl,

jax/profiler.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.profiler import (
   StepTraceAnnotation as StepTraceAnnotation,
   TraceAnnotation as TraceAnnotation,

jax/random.py

@@ -143,6 +143,8 @@ else:
   PRNGKeyArray = _PRNGKeyArray
   KeyArray = PRNGKeyArray
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
 
 from jax._src.random import (
   PRNGKey as PRNGKey,

jax/scipy/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:

jax/scipy/cluster/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JAX Authors.
+# Copyright 2022 The JAX Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,4 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax.scipy.cluster import vq as vq

jax/scipy/cluster/vq.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JAX Authors.
+# Copyright 2022 The JAX Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,4 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.cluster.vq import vq as vq

jax/scipy/fft.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.fft import (
   dct as dct,
   dctn as dctn,

jax/scipy/linalg.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.linalg import (
   block_diag as block_diag,
   cholesky as cholesky,

jax/scipy/ndimage.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.ndimage import (
   map_coordinates as map_coordinates,
 )

jax/scipy/optimize/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.optimize.minimize import (
   minimize as minimize,
   OptimizeResults as OptimizeResults,

jax/scipy/signal.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.signal import (
   convolve as convolve,
   convolve2d as convolve2d,

jax/scipy/sparse/__init__.py

@@ -12,4 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax.scipy.sparse import linalg as linalg

jax/scipy/sparse/linalg.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.sparse.linalg import (
   cg as cg,
   gmres as gmres,

jax/scipy/special.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.special import (
   betainc as betainc,
   betaln as betaln,

jax/scipy/stats/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax.scipy.stats import bernoulli as bernoulli
 from jax.scipy.stats import beta as beta
 from jax.scipy.stats import cauchy as cauchy

jax/scipy/stats/bernoulli.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.bernoulli import (
   logpmf as logpmf,
   pmf as pmf,

jax/scipy/stats/beta.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.beta import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/betabinom.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.betabinom import (
   logpmf as logpmf,
   pmf as pmf,

jax/scipy/stats/cauchy.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.cauchy import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/chi2.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.chi2 import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/dirichlet.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.dirichlet import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/expon.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.expon import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/gamma.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.gamma import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/gennorm.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.gennorm import (
   cdf as cdf,
   logpdf as logpdf,

jax/scipy/stats/geom.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.geom import (
   logpmf as logpmf,
   pmf as pmf,

jax/scipy/stats/laplace.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.laplace import (
   cdf as cdf,
   logpdf as logpdf,

jax/scipy/stats/logistic.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.logistic import (
   cdf as cdf,
   isf as isf,

jax/scipy/stats/multinomial.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
 
 from jax._src.scipy.stats.multinomial import (
   logpmf as logpmf,

jax/scipy/stats/multivariate_normal.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.multivariate_normal import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/norm.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.norm import (
   cdf as cdf,
   logcdf as logcdf,

jax/scipy/stats/pareto.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.pareto import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/poisson.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.poisson import (
   logpmf as logpmf,
   pmf as pmf,

jax/scipy/stats/t.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.t import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/truncnorm.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.truncnorm import (
   cdf as cdf,
   logcdf as logcdf,

jax/scipy/stats/uniform.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.uniform import (
   logpdf as logpdf,
   pdf as pdf,

jax/scipy/stats/vonmises.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.scipy.stats.vonmises import (
   logpdf as logpdf,
   pdf as pdf,

jax/sharding.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.sharding import (
     Sharding as Sharding,
     XLACompatibleSharding as XLACompatibleSharding,

jax/stages.py

@@ -21,6 +21,9 @@ that represent the stages of this process.
 For more, see the `AOT walkthrough <https://jax.readthedocs.io/en/latest/aot.html>`_.
 """
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.stages import (
   Compiled as Compiled,
   Lowered as Lowered,

jax/test_util.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.public_test_util import (
   check_grads as check_grads,
   check_jvp as check_jvp,

jax/tree_util.py

@@ -35,6 +35,9 @@ See the `JAX pytrees note <pytrees.html>`_
 for examples.
 """
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.tree_util import (
   Partial as Partial,
   PyTreeDef as PyTreeDef,

jax/util.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
 from jax._src.util import (
   HashableFunction as HashableFunction,
   as_hashable_function as as_hashable_function,