Marc van Zee 856d91d251 [jax2tf] Fixes a bug in testing. ...

https://github.com/google/jax/pull/11575 introduced a bug: it catches any exception raised by self.ConvertAndCompare in primitives_test.py here, but does not reraise it if the if-clause is false. This is very dangerous, since any conversion that results in an error (for instance: if an op is missing for enable_xla=False) will now pass.

I had to disable some tests that were introduced somewhere else, most likely in XLA (X64 enabled). I have filed a separate bug for this (internally).

PiperOrigin-RevId: 466269561

2022-08-09 00:34:13 -07:00

..

compilation_cache

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00

gda_serialization

Use a global module unique key. This way multiple experiments can be run and key error won't happen in WaitForBarrier.

2022-08-03 14:25:05 -07:00

jax2tf

[jax2tf] Fixes a bug in testing.

2022-08-09 00:34:13 -07:00

sparse

Make LOBPCG test plots compatible with bazel.

2022-08-04 20:05:53 -07:00

__init__.py

Delete sharded_jit

2022-07-29 08:19:52 -07:00

array.py

Make _device_assignment and _addressable_device_assignment a cached property because having a lru_cache as a decorator on the method with only self leaks memory. See the document below. In this case, its okay to make them a property rather than a method. Its also consistent with device_set.

2022-08-04 10:00:01 -07:00

callback.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

checkify.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

global_device_array.py

add an optional fastpath for api_util.shaped_abstractify

2022-07-27 15:14:37 -07:00

host_callback.py

Making sharding an arg to MHLO callback lowering

2022-08-03 11:04:14 -07:00

jet.py

[x64] make jet_test compatible with strict dtype promotion

2022-06-21 09:28:24 -07:00

maps.py

Canonicalize all shardings to OpShardingSharding throughout pjit. Places where pspec is needed, parsed_flatten_op_sharding function is used to retrieve the pspec. The major places are global_to_local and local_to_global. Rest of the changes are just threading through OpShardingSharding.

2022-08-05 12:18:10 -07:00

mesh_utils.py

[mesh_utils] Avoid relying on process-tiled device order

2022-06-06 13:30:56 -07:00

multihost_utils.py

Avoid top-level aliases of jax.tree_util.*

2022-07-07 11:41:02 -07:00

ode.py

[x64] make ode_test pass with strict dtype promotion

2022-06-17 15:21:10 -07:00

optimizers.py

[JAX] move example libraries from jax.experimental into jax.example_libraries

2021-10-19 17:30:45 -07:00

pjit.py

Add early support in pjit for single device shardings. Also lift the restriction of needing the mesh context manager when config.jax_array is enabled.

2022-08-05 22:25:25 -07:00

sharding.py

Check if axis_ctx is SPMDAxisContext because manual_axes always exists.

2022-08-06 22:36:48 -07:00

x64_context.py

[x64] make x64_context tests work with updated default dtypes

2021-12-08 13:05:40 -08:00