jax authors e03f1d4fd1 Allows for splitting the transpose of a scan into a scan and a map. ...

This is an experimental feature exposed as an extra parameter: `scan(..., _split_transpose:bool)`.

If the parameter is true then the transpose of scan generates not just 2 scans
(forward and transpose of the linearized forward), but rather 3 scans: (i)
forward (as before), (ii) transposed scan that only computes loop-carried state
required for back-propagation, but saves other intermediate gradients; (iii) a
scan (actually a map) that uses any saved activation gradients and original
residuals to compute any other gradients.

Warning: this feature is somewhat experimental and may evolve or be rolled back.
PiperOrigin-RevId: 619991098

2024-03-28 10:54:50 -07:00

..

mosaic

Allows for splitting the transpose of a scan into a scan and a map.

2024-03-28 10:54:50 -07:00

triton

Allows for splitting the transpose of a scan into a scan and a map.

2024-03-28 10:54:50 -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] Pass in compiler params via explicit compiler_params argument instead of passing via **kwargs

2024-03-06 09:16:22 -08:00

primitives.py

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

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

utils.py

Use int.bit_length() in next_power_of_2.

2024-03-14 04:31:16 -07:00