# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC

# SPDX-License-Identifier: Apache-2.0

from pathlib import Path

from models.tt_transformers.tt.lm_head import _lm_head_cache_file_name


def test_galaxy_2d_lm_head_uses_distinct_cache_key():
    cache_root = Path("weights")
    common = {
        "dummy_weights": False,
        "num_splits": 1,
        "shard": 0,
        "width": 131072,
        "mode": 0,
        "mesh_shape": (8, 4),
    }

    flat = _lm_head_cache_file_name(cache_root, galaxy_2d=False, **common)
    galaxy_2d = _lm_head_cache_file_name(cache_root, galaxy_2d=True, **common)

    assert flat.name == "output_lm_head_1_split_shard_0_131072_mode_0"
    assert galaxy_2d.name == "output_lm_head_1_split_shard_0_131072_mode_0_galaxy_2d_8x4_v1"
    assert flat != galaxy_2d


def test_dummy_lm_head_does_not_use_weight_cache():
    assert (
        _lm_head_cache_file_name(
            Path("weights"),
            dummy_weights=True,
            num_splits=1,
            shard=0,
            width=131072,
            mode=0,
            galaxy_2d=True,
            mesh_shape=(8, 4),
        )
        is None
    )