Sergei Lebedev a205c9120a pallas_call now has only one way to pass compiler_params= ...

Previously, it was possible to do

    pallas_call(..., foo=42)

and also

    pallas_call(..., compiler_params=dict(foo=42))

PiperOrigin-RevId: 623277572

2024-04-09 14:23:20 -07:00

..

mosaic

[Pallas TPU] Convert pattern_match_while_to_fori_loop to return (Jaxpr, str) rather than throw exceptions.

2024-04-09 14:04:25 -07:00

triton

The compiler_params= argument of pl.pallas_call on GPU now uses "triton" to refer to Triton-specific parameters, instead of the repetitive "triton_params"

2024-04-09 14:13:59 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

Fix some load orderings for buildifier

2024-03-27 10:28:57 -07:00

core.py

Add support for non-disjoint windows in Pallas/Mosaic

2024-02-08 02:48:43 -08:00

pallas_call.py

pallas_call now has only one way to pass compiler_params=

2024-04-09 14:23:20 -07:00

primitives.py

[pallas] Fix method name in _{load,swap}_jvp rule.

2024-03-20 11:05:22 -07:00

utils.py

[Pallas TPU] Convert pattern_match_while_to_fori_loop to return (Jaxpr, str) rather than throw exceptions.

2024-04-09 14:04:25 -07:00