xref: /aosp_15_r20/external/pytorch/test/test_ao_sparsity.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# Owner(s): ["module: unknown"]
2
3# Kernels
4from ao.sparsity.test_kernels import (  # noqa: F401
5    TestQuantizedSparseKernels,
6    TestQuantizedSparseLayers,
7)
8
9# Parametrizations
10from ao.sparsity.test_parametrization import TestFakeSparsity  # noqa: F401
11
12# Scheduler
13from ao.sparsity.test_scheduler import TestCubicScheduler, TestScheduler  # noqa: F401
14
15# Sparsifier
16from ao.sparsity.test_sparsifier import (  # noqa: F401
17    TestBaseSparsifier,
18    TestNearlyDiagonalSparsifier,
19    TestWeightNormSparsifier,
20)
21
22# Structured Pruning
23from ao.sparsity.test_structured_sparsifier import (  # noqa: F401
24    TestBaseStructuredSparsifier,
25    TestFPGMPruner,
26    TestSaliencyPruner,
27)
28
29from torch.testing._internal.common_utils import IS_ARM64, run_tests
30
31
32# Composability
33if not IS_ARM64:
34    from ao.sparsity.test_composability import (  # noqa: F401
35        TestComposability,
36        TestFxComposability,
37    )
38
39# Activation Sparsifier
40from ao.sparsity.test_activation_sparsifier import (  # noqa: F401
41    TestActivationSparsifier,
42)
43
44# Data Scheduler
45from ao.sparsity.test_data_scheduler import TestBaseDataScheduler  # noqa: F401
46
47# Data Sparsifier
48from ao.sparsity.test_data_sparsifier import (  # noqa: F401
49    TestBaseDataSparsifier,
50    TestNormDataSparsifiers,
51    TestQuantizationUtils,
52)
53
54# Utilities
55from ao.sparsity.test_sparsity_utils import TestSparsityUtilFunctions  # noqa: F401
56
57
58if __name__ == "__main__":
59    run_tests()
60