George Necula 35bfdc65e8 [shape_poly] Add some support for shape polymorphism for FFT, and tests ...

PiperOrigin-RevId: 521749241

2023-04-04 06:45:57 -07:00

..

array_serialization

Rename gda_serialization to array_serialization but keep gda_serialization around until it is included in a jax release so that OSS projects can be moved to array_serialization

2023-03-31 18:07:51 -07:00

compilation_cache

expose compiler_options on compile()

2023-03-30 17:14:26 -07:00

gda_serialization

Add Deprecation warning if gda_serialization is imported

2023-03-31 21:28:07 -07:00

jax2tf

[shape_poly] Add some support for shape polymorphism for FFT, and tests

2023-04-04 06:45:57 -07:00

sparse

Merge pull request #15243 from jakevdp:coo-sort-warning

2023-03-27 16:50:24 -07:00

__init__.py

Change jax.ad, jax.xla, jax.pxla to point to the shims instead of the internal modules.

2023-02-09 13:31:40 -08:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

host_callback.py

Bump minimum jaxlib version from 0.4.6 to 0.4.7.

2023-03-28 13:43:01 -07:00

jet.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

maps.py

Remove _PositionalSemantics class since it is not used anymore because jax.Array always has GLOBAL semantics

2023-03-17 13:30:04 -07:00

mesh_utils.py

Enable more mesh shape assignment

2023-03-31 09:36:16 -07:00

multihost_utils.py

Add deprecation warning for FROM_GDA usage since that argument is not required anymore.

2023-03-27 11:33:11 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

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

pjit.py

Delete unused functions from dispatch.py and pjit.py

2023-03-30 13:38:44 -07:00

rnn.py

handle seq_lengths in lstm_ref

2023-04-03 22:22:54 +00:00

serialize_executable.py

Redefine compile_and_serialize as serialize(lowered.compile()).

2023-03-22 17:23:19 -07:00

shard_map.py

[shard-map] handle closed-over vmap tracers

2023-03-30 16:43:40 -07:00

x64_context.py

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

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