Add fully configurable activation checkpointing (#951)
* add fully configurable activation checkpointing
* fix format
* fix format
* add docstring to activation_checkpointing_fn
* add block id range option in act ckpt
* resolve conflict
* add a check for blocks ids overlap in mapping
* fix typo
* update docstring
* refactor
* fix test
* Apply suggestions from code review
Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
* address comments
* add build mapping as a helper func
* fix format
---------
Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
