George Necula 6e09d4a71c [shape_poly] Change lowering for shape polymorphism to simplify shape refinement ...

We insert a ConvertOp as the only use of an input argument in a shape polymorphic
`main` function. This helps the downstream shape refinement because it will set the type
of input arguments to static shapes, and this can invalidate the
module if the argument appears as the result of a function, or if
it appears as the input to a custom_call with output_operand_alias
attribute.
See b/287386268.

2023-06-15 11:58:02 +03:00

..

array_serialization

Fix get_tensorstore_spec for GCS paths if ocdbt is enabled

2023-06-07 16:47:35 -07: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

[shape_poly] Change lowering for shape polymorphism to simplify shape refinement

2023-06-15 11:58:02 +03:00

sparse

Make dimensions argument of bcoo_reshape optional.

2023-05-09 10:38:18 -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

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

jet.py

Split sharding_impls into its own Bazel target.

2023-04-10 10:15:58 -07:00

maps.py

Copybara import of the project:

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

mesh_utils.py

Change uses of np.product to np.prod

2023-06-02 03:57:30 -07:00

multihost_utils.py

Make inspect_sharding_lowering_rule work only with HloSharding in it's callback. Also remove get_replicated_op_sharding since it is not needed anymore.

2023-06-07 14:39:52 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

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

pjit.py

Copybara import of the project:

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

rnn.py

Make rnn_bwd_abstract_eval backwards compatible by guarding it agains the jaxlib version

2023-05-03 19:28:42 -07:00

serialize_executable.py

Stop passing CompileOptions when deserializing.

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

shard_map.py

Merge pull request #16072 from axch:stacked-piles

2023-06-08 11:22:30 -07:00

topologies.py

Default None-topology platform to TPU.

2023-05-11 10:41:45 -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