ColossalAI/tests/test_fx/test_codegen
..
test_activation_checkpoint_codegen.py