George Necula ec8b855fa1 [shape_poly] Add a polymorphic shape refinement MLIR pass accessible to JAX Python. ...

At the moment we can run the StableHLO module lowered by jax2tf
with polymorphic shapes only with jax2tf, because the tf.XlaCallModule op has the
necessary shape refinement logic (which is necessary to legalize
the StableHLO module with dynamic shapes to MHLO). Here we
expose the shape refinement MLIR transformation to JAX Python.

For now this is used only in a test in jax_export_test.py.

PiperOrigin-RevId: 537485288

2023-06-02 21:49:20 -07:00

..

_src

Merge pull request #16231 from jakevdp:product

2023-06-02 18:13:56 -07:00

example_libraries

Link to JAXopt in example_libraries.optimizers

2023-01-25 10:54:31 -08:00

experimental

[shape_poly] Add a polymorphic shape refinement MLIR pass accessible to JAX Python.

2023-06-02 21:49:20 -07:00

extend

initiate jax.extend via docs and top-level module set-up

2023-05-15 15:47:06 -07:00

image

Document why 'import name as name' is used

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

interpreters

Deprecate jax.interpreters.xla.Buffer, device_put, xla_call_p

2023-05-28 07:15:34 -07:00

lax

deprecate jax.lax.prod

2023-05-23 17:33:50 -07:00

lib

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

nn

Add test framework for module attribute

2023-04-21 13:20:16 -07:00

numpy

Deprecate jax.numpy functions alltrue, sometrue, product, cumproduct

2023-06-02 04:10:46 -07:00

ops

Document why 'import name as name' is used

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

scipy

feat (scipy.special): Add a xla version of scipy.special.gamma function

2023-04-18 21:10:22 +02:00

tools

Migrate exec_tools back to tools.

2023-05-23 14:00:34 -07:00

__init__.py

add custom_jvp for jax.nn.softmax

2023-04-22 11:28:03 -07:00

abstract_arrays.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

ad_checkpoint.py

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

2022-09-22 12:27:19 -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

Break cycle between _src/core.py and _src/dtypes.py.

2023-05-17 07:58:59 -07: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

Cleanup: remove unused imports in private modules

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

core.py

Break cycle between _src/core.py and _src/dtypes.py.

2023-05-17 07:58:59 -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

Remove jax.dtypes._jax_types

2023-02-01 12:49:06 -08:00

errors.py

Document why 'import name as name' is used

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

flatten_util.py

ravel_pytree now accepted by static linters

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

jaxpr_util.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

linear_util.py

jax.linear_util: remove unnecessary exported names

2022-12-21 18:52:01 -08:00

monitoring.py

Move contents of jax.monitoring to jax._src.monitoring.

2023-02-08 15:03:22 -08:00

prng.py

Document why 'import name as name' is used

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

profiler.py

Document why 'import name as name' is used

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

py.typed

Improve type annotations

2021-01-13 10:26:35 -05:00

random.py

add new random key constructor

2023-05-22 11:35:10 -07:00

sharding.py

Move PartitionSpec into its own file (jax/_src/partition_spec.py).

2023-04-06 11:43:25 -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

Stop importing old tree_util APIs conveniently and set explicit time for removal.

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

typing.py

jax.typing: recommend instance check in Python 3.10 or newer

2023-03-27 10:01:28 -07:00

util.py

Document why 'import name as name' is used

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

version.py

Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332

2023-06-01 09:42:55 -07:00