Matthew Johnson e3a92d52ba prepare to switch to new remat ...

This commit involves a few things, which are all united in being about landing
the new remat (aka new checkpoint) implementation:
  * add benchmarks for new remat eager performance, and some caching to make those
    benchmarks fast
  * warn when the old-remat-exclusive `concrete` feature is used, with an
    actionable message pointing to the new recommended approach involving static_argnums
  * add the static_argnums parameter to both new and old remt
  * update docstrings (and de-duplicate them to)
  * add new tests, especially around caching and errors/warnings

2022-08-04 12:25:03 -07:00

..

api_benchmark.py

prepare to switch to new remat

2022-08-04 12:25:03 -07:00

benchmark.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

gda_benchmark.py

Fix the gda_xla_sharding_match benchmark which was regressing. This was happening because that function was executed from top to bottom a couple of times and each time a new mesh object was created violating the already created cache which doesn't happen in real life.

2022-07-27 23:08:55 -07:00

pmap_benchmark.py

Revert import path changes to examples/ and benchmarks/

2021-01-20 17:35:55 -08:00