Adam Paszke 8feab68209 [Mosaic GPU] Remove the unnecessary scratch space operand ...

And clean up the C++ dispatch code. We don't use HBM scratch anymore
since we pass TMA descriptors as kernel arguments.

PiperOrigin-RevId: 671327420

2024-09-05 04:57:52 -07:00

..

mosaic

Added pl.CompilerParams subclass for Mosaic GPU

2024-09-04 12:48:34 -07:00

mosaic_gpu

[Mosaic GPU] Remove the unnecessary scratch space operand

2024-09-05 04:57:52 -07:00

triton

Add TritonCompilerParams for specifying compiler arguments instead of a dict.

2024-09-04 13:32:25 -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

Added pl.CompilerParams subclass for Mosaic GPU

2024-09-04 12:48:34 -07:00

pallas_call.py

Merge pull request #23442 from superbobry:pyright

2024-09-05 02:46:26 -07:00

primitives.py

Fixed some type errors under pyright

2024-09-05 09:56:38 +01:00

utils.py

[Pallas GPU] Add 32-bit lowering rule for lax.erf_inv

2024-08-28 17:59:15 -07:00