Peter Hawkins 4f805c2d8f [JAX] Change jax.test_util utilities to have identical tolerances on all platforms. ...

In cases where this causes TPU tests to fail, relax test tolerances in the test cases themselves.

TPUs are less precise only for specific operations, notably matrix multiplication (for which usually enabling higher-precision matrix multiplication is the right choice if precision is needed), and certain special functions (e.g., log/exp/pow).

The net effect of this change is mostly to tighten up many test tolerances on TPU.

PiperOrigin-RevId: 562953488

2023-09-05 18:48:55 -07:00

..

array_serialization

Abstract the array_serialization error message to a global variable so that it can be overridden.

2023-08-29 21:46:26 -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

jax2tf

[JAX] Change jax.test_util utilities to have identical tolerances on all platforms.

2023-09-05 18:48:55 -07:00

mosaic

Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic

2023-07-20 18:28:51 -07:00

pallas

Add segment_ids support to pallas flash attention on GPU.

2023-08-29 13:59:18 -07:00

sparse

Lower jax.numpy matmul functions to mixed-precision dot_general

2023-09-05 08:37:51 -07: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

Pad generated sharding specs with None up to ndims to simplify comparing dims

2023-08-10 17:02:31 -07:00

host_callback.py

[callback] Enable device_index support in terms of callback sharding support.

2023-08-31 22:31:35 -07:00

jet.py

Use lower-case PEP 585 names for types.

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

maps.py

Copybara import of the project:

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

mesh_utils.py

Allow device mesh handler to return None and use the default logic

2023-08-07 12:50:43 -07:00

multihost_utils.py

Add docs for is_fully_addressable to jax.Array and remove GDA from process_allgather docs and clarify it a bit more.

2023-08-20 18:58:22 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

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

pjit.py

Prune some exports from jax.experimental.pjit.

2023-08-01 13:27:17 -07: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

Use lower-case PEP 585 names for types.

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

serialize_executable.py

Stop passing CompileOptions when deserializing.

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

shard_map.py

rolling forward shard_map transpose fixes

2023-08-31 17:31:21 -07:00

topologies.py

Apply pyupgrade --py39-plus.

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

x64_context.py

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

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