mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 21:46:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/experimental
History
jax authors 1ceddfc98a Merge pull request #16710 from gnecula:poly_max0 
PiperOrigin-RevId: 549515427
2023-07-19 21:40:17 -07:00
..
array_serialization
Improve error message and the relevant test.
2023-07-14 01:45:43 -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
Merge pull request #16710 from gnecula:poly_max0
2023-07-19 21:40:17 -07:00
sparse
[sparse] support preferred_element_type in dot_general
2023-07-14 18:23:34 -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
[JAX] Add support for multiple pytree registries.
2023-07-19 06:48:21 -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
Improve error message to point the way to Megacore.
2023-07-11 08:16:33 -07:00
multihost_utils.py
Raise an error if a user passes None to host_local_array_to_global_array or global_array_to_host_local_array
2023-06-26 18:15:43 -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
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
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
topologies.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -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