Module neuralops.tests.test_admm_pruner_distill