Jake VanderPlas 3cce03554f [sparse] remove deprecated _dedupe() method

2022-08-01 16:59:25 -07:00

..

compilation_cache

Update minimum jaxlib version to 0.3.14.

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

gda_serialization

Log the key used for barrier

2022-07-27 20:25:40 -07:00

jax2tf

Merge pull request #11653 from sharadmv:debugging-docs

2022-07-28 20:26:25 -07:00

sparse

[sparse] remove deprecated _dedupe() method

2022-08-01 16:59:25 -07:00

__init__.py

Delete sharded_jit

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

array.py

add an optional fastpath for api_util.shaped_abstractify

2022-07-27 15:14:37 -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

fix ad_checkpoint.checkpoint caching issue

2022-07-29 19:59:28 -07:00

jet.py

[x64] make jet_test compatible with strict dtype promotion

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

maps.py

add custom-policy partial eval and dce rules for pmap

2022-07-28 21:13:25 -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

Make is_compatible_aval an optional method which sharding interfaces can implement to raise a more meaningful error. Otherwise lower to opsharding and catch the error if it fails.

2022-07-29 13:37:15 -07:00

sharding.py

Repr the op_sharding for printing.

2022-07-30 15:25:26 -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