George Necula 3601b25899 Move multi_platform_export_test.py out of jax2tf. ...

This test is now independent of jax2tf. Move it out and rename it export_harnesses_multi_platform_test.py.

We disable the test in GitHub CI, because it is very large, pending
some changes to ensure it parallelizes well. The test is still
running in internal CI. This is matching the current behavior, since
jax2tf tests are only run internally.

PiperOrigin-RevId: 583603863

2023-11-18 02:52:44 -08:00

..

_src

Move multi_platform_export_test.py out of jax2tf.

2023-11-18 02:52:44 -08:00

example_libraries

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

experimental

Move multi_platform_export_test.py out of jax2tf.

2023-11-18 02:52:44 -08:00

extend

[jex] replace extend.random.PRNGImpl with extend.random.define_prng_impl

2023-10-31 17:21:54 -07:00

image

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

interpreters

jax.interpreters.ad: deprecate inadvertently exported symbols

2023-10-16 11:22:37 -07:00

lax

Add a lax.platform_dependent API for writing platform-dependent code.

2023-11-02 14:31:38 +01:00

lib

Switch from flake8 to Ruff

2023-11-15 22:35:52 -05:00

nn

Add nn.squareplus.

2023-11-14 23:52:41 -05:00

numpy

Annotate jax.numpy.atleast_xd functions

2023-11-09 19:21:35 -05:00

ops

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

scipy

implementation of poch and hyp1f1

2023-11-15 20:01:00 +01:00

tools

[PJRT C API] Change build wheel script to build a separate package for cuda kernels.

2023-11-06 09:13:44 -08:00

__init__.py

add jax.threefry_partitionable context manager

2023-10-31 13:45:55 -07:00

ad_checkpoint.py

Add memories support to remat.

2023-09-12 20:50:05 -07:00

api_util.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

BUILD

Move back_compat_test_util.py to jax._src.internal_test_util.

2023-11-17 02:05:30 -08:00

cloud_tpu_init.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

collect_profile.py

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

config.py

Deprecated accessing config via the jax.config submodule

2023-10-27 20:08:10 +01:00

core.py

jax.core: point deprecation to jax.extend

2023-10-13 12:49:05 -07:00

custom_batching.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

custom_derivatives.py

custom_vjp symbolic zeros support, take two

2023-04-05 11:17:05 -07:00

custom_transpose.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

debug.py

Explicitly re-export all names from jax.debug

2023-04-07 07:24:59 -07:00

distributed.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

dlpack.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

dtypes.py

Copybara import of the project:

2023-07-24 14:38:20 -07:00

errors.py

Tracebacks no longer have JAX-internal frames prepended by default

2023-08-03 11:38:38 -07:00

flatten_util.py

ravel_pytree now accepted by static linters

2023-03-07 23:35:48 +00:00

linear_util.py

JEX: move jax.linear_util to jax.extend.linear_util

2023-08-30 18:32:12 -07:00

monitoring.py

Allow duration event listeners to take extra keyword arguments.

2023-10-04 13:15:49 -07:00

prng.py

Deprecate the contents of jax.prng

2023-08-30 15:13:32 -07:00

profiler.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

py.typed

…

random.py

Deprecate jax.random.shuffle

2023-11-06 12:21:56 -08:00

sharding.py

Delete OpShardingSharding export since it has been 3 months since it was deprecated. Also remove deprecation warnings for MeshPspecSharding.

2023-06-08 13:53:38 -07:00

stages.py

add CompilerOptions type alias for AOT compiler options

2023-04-12 16:49:13 -07:00

test_util.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

tree_util.py

fix typo in deprecation message

2023-10-03 14:04:09 -07:00

typing.py

Define jax.typing.DTypeLike

2023-10-10 08:46:36 -07:00

util.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

version.py

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00