George Necula d25bcac93d [shape_poly] Add better support for division, and working with strides ...

Previously, division was only supported in certain situation, and this
led to errors, e.g., when using strides. Now we generalize the polynomials
to also include "floordiv(E, E)" and "mod(E, E)" as atoms, in addition
to dimension variables. A symbolic dimension is now a sum of products
of atoms. (We also changed the documentation to use symbolic dimension
instead of dimension polynomials).

2023-01-25 07:37:54 -08:00

..

compilation_cache

Add {SerializeAs,ParseFrom}String to xe.CompileOptions.

2023-01-03 09:15:34 -08:00

gda_serialization

Add gcs support to serialization for internal tensor spec.

2023-01-18 16:43:03 -08:00

jax2tf

[shape_poly] Add better support for division, and working with strides

2023-01-25 07:37:54 -08:00

sparse

Merge pull request #14091 from jakevdp:bcoo-methods

2023-01-23 10:11:14 -08:00

__init__.py

Add IO callback

2023-01-17 13:55:05 -08:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

cleanup: remove several unused imports across the package

2022-12-28 12:49:17 -08:00

global_device_array.py

Delete the pjit 101 tutorial which is subsumed by https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html

2023-01-20 20:35:47 -08:00

host_callback.py

Raise an error when attempting to mutate Jaxpr objects

2023-01-23 09:37:58 -08:00

jet.py

Add pjit support in jax.experimental.jet

2023-01-23 15:51:47 -08:00

maps.py

Make xmap use dispatch.sharded_lowering as dispatch.lower_xla_callable is deprecated.

2023-01-16 09:42:40 -08:00

mesh_utils.py

Delete the pjit 101 tutorial which is subsumed by https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html

2023-01-20 20:35:47 -08:00

multihost_utils.py

Add typehints and point to the correct endpoint of Mesh and PartitionSpec in the args section.

2022-12-05 09:49:18 -08:00

ode.py

fix transcription error for initial_step_size (thanks @AlbertMitjans)

2023-01-14 21:10:44 -08:00

pjit.py

Move pjit.py to jax/_src in preparation for merging the jit and pjit frontend APIs

2022-12-16 13:07:15 -08:00

rnn.py

Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.

2022-11-28 14:31:48 -08:00

serialize_executable.py

Refactor create_cpp_call to be a method on MeshExecutable

2022-11-17 18:05:10 -08:00

x64_context.py

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

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