George Necula 4fbf50dd60 [shape_poly] Copy many of the jax2tf/shape_poly_test to live outside of jax2tf. ...

Shape polymorphism is now usable independently of jax2tf, and it deserves to have its tests independent of jax2tf. I started by branching jax2tf/tests/shape_poly_test.py into tests/shape_poly_test.py, followed by removing from the latter the tests and helper functions that do not make sense outside of jax2tf.

For now we leave the existing tests in jax2tf, because some of those tests exercise
other code paths. In the process of adding these tests we found two bugs (fixed separately in https://github.com/google/jax/pull/18516 and https://github.com/google/jax/pull/18515).

Since we now run these tests in GitHub and Kokoro, this has revealed a couple
of bugs in the tests, which we fix here both in the jax2tf/tests/shape_poly_test.py and the copy tests/shape_poly_test.py.

PiperOrigin-RevId: 583816243

2023-11-19 09:00:04 -08:00

..

array_api

Add jax.experimental.array_api interface

2023-11-16 14:21:04 -08:00

array_serialization

Add mitigation techniques in the error message when a barrier timeout occurs.

2023-10-27 08:57:39 -07:00

compilation_cache

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

export

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08:00

jax2tf

[shape_poly] Copy many of the jax2tf/shape_poly_test to live outside of jax2tf.

2023-11-19 09:00:04 -08:00

mosaic

Increase minimum jaxlib version to 0.4.19.

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

pallas

Avoids some unnecessary dtype casts in Pallas attention to make fwd pass faster. Use num_warps=8 to make bwd faster. Minor fixes to make code more readable.

2023-11-13 15:46:46 -08:00

sparse

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08:00

__init__.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Copybara import of the project:

2023-11-06 17:01:02 -08:00

host_callback.py

Cleanup the handling of single- and multi-platform lowering in ModuleContext

2023-10-25 10:40:41 -07:00

jet.py

Use lower-case PEP 585 names for types.

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

layout.py

Remove DefaultLayout and make None same as DefaultLayout

2023-11-16 18:01:27 -08:00

maps.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

mesh_utils.py

Switch from flake8 to Ruff

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

multihost_utils.py

Switch from flake8 to Ruff

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

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

pjit.py

Switch from flake8 to Ruff

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

profiler.py

Add get_profiled_instructions_proto API in jax.experimental.profiler to get profiled instructions proto.

2023-06-29 21:02:05 -07:00

rnn.py

Increase minimum jaxlib version to 0.4.19.

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

serialize_executable.py

Stop passing CompileOptions when deserializing.

2023-05-10 16:22:54 -07:00

shard_map.py

factor out subs_list and subs_list2

2023-10-19 09:32:44 -07:00

topologies.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

x64_context.py

MAINT Do not import the config object in JAX internals

2023-10-18 10:55:13 +01:00