pytorch

1
# Owner(s): ["module: unknown"]
2

3
# Kernels
4
from ao.sparsity.test_kernels import (  # noqa: F401
5
    TestQuantizedSparseKernels,
6
    TestQuantizedSparseLayers,
7
)
8

9
# Parametrizations
10
from ao.sparsity.test_parametrization import TestFakeSparsity  # noqa: F401
11

12
# Scheduler
13
from ao.sparsity.test_scheduler import TestCubicScheduler, TestScheduler  # noqa: F401
14

15
# Sparsifier
16
from ao.sparsity.test_sparsifier import (  # noqa: F401
17
    TestBaseSparsifier,
18
    TestNearlyDiagonalSparsifier,
19
    TestWeightNormSparsifier,
20
)
21

22
# Structured Pruning
23
from ao.sparsity.test_structured_sparsifier import (  # noqa: F401
24
    TestBaseStructuredSparsifier,
25
    TestFPGMPruner,
26
    TestSaliencyPruner,
27
)
28

29
from torch.testing._internal.common_utils import IS_ARM64, run_tests
30

31

32
# Composability
33
if not IS_ARM64:
34
    from ao.sparsity.test_composability import (  # noqa: F401
35
        TestComposability,
36
        TestFxComposability,
37
    )
38

39
# Activation Sparsifier
40
from ao.sparsity.test_activation_sparsifier import (  # noqa: F401
41
    TestActivationSparsifier,
42
)
43

44
# Data Scheduler
45
from ao.sparsity.test_data_scheduler import TestBaseDataScheduler  # noqa: F401
46

47
# Data Sparsifier
48
from ao.sparsity.test_data_sparsifier import (  # noqa: F401
49
    TestBaseDataSparsifier,
50
    TestNormDataSparsifiers,
51
    TestQuantizationUtils,
52
)
53

54
# Utilities
55
from ao.sparsity.test_sparsity_utils import TestSparsityUtilFunctions  # noqa: F401
56

57

58
if __name__ == "__main__":
59
    run_tests()
60