# Owner(s): ["module: sparse"]
import logging

# Kernels
from ao.sparsity.test_kernels import (  # noqa: F401
    TestQuantizedSparseKernels,
    TestQuantizedSparseLayers,
)

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

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

# Sparsifier
from ao.sparsity.test_sparsifier import (  # noqa: F401
    TestBaseSparsifier,
    TestNearlyDiagonalSparsifier,
    TestWeightNormSparsifier,
)

# Structured Pruning
from ao.sparsity.test_structured_sparsifier import (  # noqa: F401
    TestBaseStructuredSparsifier,
    TestFPGMPruner,
    TestSaliencyPruner,
)

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


# Composability
if not IS_ARM64:
    from ao.sparsity.test_composability import (  # noqa: F401
        TestComposability,
        TestFxComposability,
    )

# Activation Sparsifier
from ao.sparsity.test_activation_sparsifier import (  # noqa: F401
    TestActivationSparsifier,
)

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

# Data Sparsifier
from ao.sparsity.test_data_sparsifier import (  # noqa: F401
    TestBaseDataSparsifier,
    TestNormDataSparsifiers,
    TestQuantizationUtils,
)

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


if __name__ == "__main__":
    logging.basicConfig(
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        level=logging.INFO,
    )

    run_tests()
