X_SSL_Net/tests/test_csm_triton_cpu.py

from __future__ import annotations

import sys
from pathlib import Path

import torch

ROOT_DIR = Path(__file__).resolve().parents[1]
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))


def test_cross_scan_fn_force_torch_supports_cpu() -> None:
    from lib.modules.lib_mamba.csm_triton import cross_scan_fn

    x = torch.randn(1, 8, 4, 4)
    y = cross_scan_fn(x, force_torch=True)

    assert y.shape == (1, 4, 8, 16)


def test_cross_merge_fn_force_torch_supports_cpu() -> None:
    from lib.modules.lib_mamba.csm_triton import cross_merge_fn

    y = torch.randn(1, 4, 8, 4, 4)
    x = cross_merge_fn(y, force_torch=True)

    assert x.shape == (1, 8, 16)