Michael Hudgins d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax ...

PiperOrigin-RevId: 676843138

2024-09-20 07:52:33 -07:00

..

mosaic

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

mosaic_gpu

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

triton

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

Clean up BUILD files.

2024-08-26 09:11:17 -07:00

core.py

[Pallas] Minor cleanup of memory spaces. Also add ANY as a general memory space

2024-09-19 19:08:18 -07:00

pallas_call.py

[Pallas TPU] Fix some issues introduced by the recent changes

2024-09-19 04:39:59 -07:00

primitives.py

[Pallas TPU] Add support for passing in and returning semaphores

2024-09-18 20:39:43 -07:00

utils.py

[Pallas TPU] Add lowering for 64 bit

2024-09-19 16:42:45 +01:00