ColossalAI/tests/test_layers/test_2d
..
checks_2d
test_2d.py