Tianjian Lu 020849076c [linalg] Add tpu svd lowering rule. ...

PiperOrigin-RevId: 445533767

2022-04-29 16:43:53 -07:00

..

compilation_cache

Increase minimum jaxlib version to 0.3.7.

2022-04-18 08:09:50 -04:00

gda_serialization

Add a dtypes option to cast host arrays when reloading from TS.

2022-04-22 18:00:27 -07:00

jax2tf

[linalg] Add tpu svd lowering rule.

2022-04-29 16:43:53 -07:00

sparse

[sparse] implement sparse rule for lax.concatenate_p

2022-04-28 10:59:11 -07:00

__init__.py

CI: fix flake8 ignore declarations

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

callback.py

Switch internal users of jax.util.partial to use functools.partial.

2021-09-13 21:09:58 -04:00

checkify.py

CI: fix flake8 ignore declarations

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

global_device_array.py

Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint).

2022-04-06 15:19:19 -07:00

host_callback.py

Add an emit_python_callback helper function

2022-04-26 12:19:50 -07:00

jet.py

remove _reduce_min from public jax.lax module

2022-03-08 16:34:26 -08:00

loops.py

Deprecate jax.experimental.loops, step 2.

2022-04-14 12:38:55 -07:00

maps.py

reviewer comments

2022-04-29 09:09:07 -07:00

mesh_utils.py

[mesh_utils] Add device/slice count checks

2022-04-20 13:30:01 -07:00

multihost_utils.py

Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint).

2022-04-06 15:19:19 -07:00

ode.py

Fix a bug that promoted t to a complex in odeint, and modify a test so it would have caught it

2022-01-07 12:10:35 -05:00

optimizers.py

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

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

pjit.py

[remove-units] prevent pjit partial eval from dealing with units

2022-04-27 20:07:20 -07:00

stax.py

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

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

x64_context.py

[x64] make x64_context tests work with updated default dtypes

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