Sharad Vikram d872812a35 [Pallas] Upstream pallas to JAX ...

PiperOrigin-RevId: 552963029

2023-08-01 16:43:13 -07:00

..

array_serialization

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

compilation_cache

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

jax2tf

Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54

2023-08-01 08:53:28 -07:00

mosaic

Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic

2023-07-20 18:28:51 -07:00

pallas

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

sparse

Previously, using sparse.todense on a BCSR matrix with sparse.sparsify would raise NotImplementedError: sparse rule for todense is not implemented. By adding the sparse rule, it will resolve this issue.

2023-07-26 13:01:02 -07:00

__init__.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Return PositionalSharding instead of GSPMDSharding in custom_partitioning when mesh is not defined

2023-06-12 11:52:28 -07:00

host_callback.py

Update flags to use the ABSL typed flag API.

2023-07-27 12:15:58 -07:00

jet.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

maps.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

mesh_utils.py

Don't depend on jax in mesh_utils to remove circular dependency.

2023-08-01 07:45:42 -07:00

multihost_utils.py

[JAX] Remove code that sets or tests --jax_coordination_service.

2023-07-28 13:13:13 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

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

pjit.py

Prune some exports from jax.experimental.pjit.

2023-08-01 13:27:17 -07:00

profiler.py

Add get_profiled_instructions_proto API in jax.experimental.profiler to get profiled instructions proto.

2023-06-29 21:02:05 -07:00

rnn.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

serialize_executable.py

Stop passing CompileOptions when deserializing.

2023-05-10 16:22:54 -07:00

shard_map.py

[shmap-partial-auto] start adding partial auto / partial manual to shmap

2023-08-01 14:51:07 -07:00

topologies.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

x64_context.py

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

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