Sharad Vikram 86023f55ea [Pallas TPU] Add DMA descriptor abstraction for constructing but not starting DMAs ...

PiperOrigin-RevId: 574210634

2023-10-17 11:20:05 -07:00

..

mosaic

[Pallas TPU] Add DMA descriptor abstraction for constructing but not starting DMAs

2023-10-17 11:20:05 -07:00

triton

[pallas:gpu] Simplify broadcast_to, min, max lowering.

2023-10-17 11:00:50 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

Load Python rules from rules_python

2023-08-24 10:22:57 -07:00

core.py

[Pallas] Add support for pytrees in scalar prefetch

2023-10-03 11:52:34 -07:00

indexing.py

[pallas] Simplify Slice.from_slice code and add check for Slice.size.

2023-10-13 01:36:28 -07:00

pallas_call.py

[Pallas] Add async_copy_to and async_remote_copy_to for doing DMAs.

2023-09-14 14:32:08 -07:00

primitives.py

[pallas] dot fixes.

2023-10-16 12:35:43 -07:00

utils.py

[Pallas] Add support for lax.fori_loop to Mosaic lowering

2023-08-04 13:43:43 -07:00