code/models/common/tests/test_lazy_weight.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Unit tests for LazyWeight module.
"""

import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
import torch

import ttnn
from models.common.modules.lazy_weight import (
    LazyWeight,
    _auto_pad_for_sharded_tiles,
    _from_torch_and_dump,
    resolve_lazy_weight,
)

# ============================================================================
# Fixtures
# ============================================================================


@pytest.fixture
def mock_tensor():
    """Create a mock torch tensor for testing."""
    return torch.randn(1, 32, 128)


@pytest.fixture
def mock_tensor_2d():
    """Create a 2D mock torch tensor for testing."""
    return torch.randn(64, 128)


@pytest.fixture
def mock_mesh_device():
    """Create a mock MeshDevice for unit tests."""
    device = MagicMock(spec=ttnn.MeshDevice)
    device.get_num_devices.return_value = 8
    device.id.return_value = "test_device_0"
    return device


@pytest.fixture
def mock_memory_config():
    """Create a mock MemoryConfig."""
    return ttnn.DRAM_MEMORY_CONFIG


@pytest.fixture
def mock_mesh_mapper_config():
    """Create a mock MeshMapperConfig for sharding."""
    return ttnn.MeshMapperConfig(
        placements=[ttnn.PlacementShard(-1)],
        mesh_shape_override=ttnn.MeshShape([8]),
    )


# ============================================================================
# Test LazyWeight Construction
# ============================================================================


class TestLazyWeightConstruction:
    """Tests for LazyWeight dataclass construction and __post_init__."""

    def test_construction_with_tensor(self, mock_tensor):
        """Test basic construction with a tensor."""
        lw = LazyWeight(source=mock_tensor)
        assert lw.source is mock_tensor
        assert lw.dtype == ttnn.bfloat16
        assert lw.pad_value == 0.0
        assert lw.device is None
        assert lw.mesh_mapper_config is None
        assert lw.memory_config is None
        assert lw.layout is None
        assert lw._value is None

    def test_construction_with_all_parameters(
        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
    ):
        """Test construction with all parameters specified."""
        cache_dir = Path("/tmp/cache")
        lw = LazyWeight(
            source=mock_tensor,
            cache_dir_weight_name=(cache_dir, "test_weight"),
            pad_value=1.0,
            dtype=ttnn.bfloat4_b,
            device=mock_mesh_device,
            mesh_mapper_config=mock_mesh_mapper_config,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        assert lw.source is mock_tensor
        assert lw.cache_dir_weight_name == (cache_dir, "test_weight")
        assert lw.pad_value == 1.0
        assert lw.dtype == ttnn.bfloat4_b
        assert lw.device is mock_mesh_device
        assert lw.mesh_mapper_config is mock_mesh_mapper_config
        assert lw.memory_config is mock_memory_config
        assert lw.layout == ttnn.TILE_LAYOUT

    def test_post_init_validates_shape(self):
        """Test that __post_init__ validates tensor shape."""
        mock_source = MagicMock()
        mock_source.shape = None

        with pytest.raises(AssertionError, match="source must have a shape"):
            LazyWeight(source=mock_source)

    def test_post_init_validates_empty_shape(self):
        """Test that __post_init__ rejects empty shape."""
        mock_source = MagicMock()
        mock_source.shape = ()

        with pytest.raises(AssertionError, match="source must have a shape"):
            LazyWeight(source=mock_source)


# ============================================================================
# Test is_resolved Method
# ============================================================================


class TestIsResolved:
    """Tests for LazyWeight.is_resolved() method."""

    def test_not_resolved_missing_device(self, mock_tensor, mock_memory_config, mock_mesh_mapper_config):
        """Test is_resolved returns False when device is missing."""
        lw = LazyWeight(
            source=mock_tensor,
            memory_config=mock_memory_config,
            mesh_mapper_config=mock_mesh_mapper_config,
            layout=ttnn.TILE_LAYOUT,
        )
        assert not lw.is_resolved()

    def test_not_resolved_missing_layout(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test is_resolved returns False when layout is missing."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
        )
        assert not lw.is_resolved()

    def test_not_resolved_missing_memory_config(self, mock_tensor, mock_mesh_device):
        """Test is_resolved returns False when memory_config is missing."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            layout=ttnn.TILE_LAYOUT,
        )
        assert not lw.is_resolved()

    def test_is_resolved_with_all_required_fields(
        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
    ):
        """Test is_resolved returns True when all required fields are set."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            mesh_mapper_config=mock_mesh_mapper_config,
            layout=ttnn.TILE_LAYOUT,
        )
        assert lw.is_resolved()

    def test_is_resolved_with_none_mesh_mapper_config(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test is_resolved returns False when mesh_mapper_config is None (replicated case requires it to be explicitly set)."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper_config=None,  # None means not resolved per is_resolved check
        )
        # mesh_mapper_config is in required_fields so None means not resolved
        assert not lw.is_resolved()


# ============================================================================
# Test _get_fingerprint Method
# ============================================================================


class TestGetFingerprint:
    """Tests for LazyWeight._get_fingerprint() method."""

    def test_fingerprint_includes_source_shape(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint includes source shape."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        fingerprint = lw._get_fingerprint()
        assert "srcshape_1_32_128" in fingerprint

    def test_fingerprint_includes_non_default_memory_config(self, mock_tensor, mock_mesh_device):
        """Test that fingerprint includes memory_config hash when non-default."""
        # Use L1_MEMORY_CONFIG which is different from DRAM_MEMORY_CONFIG
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=ttnn.L1_MEMORY_CONFIG,
            layout=ttnn.TILE_LAYOUT,
        )
        fingerprint = lw._get_fingerprint()
        assert "memcfg_" in fingerprint

    def test_fingerprint_includes_dtype(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint includes dtype."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            dtype=ttnn.bfloat16,
        )
        fingerprint = lw._get_fingerprint()
        assert "dtype_BFLOAT16" in fingerprint

    def test_fingerprint_includes_layout(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint includes layout."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        fingerprint = lw._get_fingerprint()
        assert "layout_TILE" in fingerprint

    def test_fingerprint_includes_pad_value_when_non_default(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint includes pad_value when non-default."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            pad_value=1.5,
        )
        fingerprint = lw._get_fingerprint()
        assert "pad_1.5" in fingerprint

    def test_fingerprint_excludes_pad_value_when_default(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint excludes pad_value when default (0.0)."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            pad_value=0.0,
        )
        fingerprint = lw._get_fingerprint()
        assert "pad_" not in fingerprint

    def test_fingerprint_includes_device_id(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint includes device id."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        fingerprint = lw._get_fingerprint()
        assert "device_" in fingerprint

    def test_fingerprint_device_without_id_method(self, mock_tensor, mock_memory_config):
        """Test fingerprint uses 'single' when device has no id() method."""
        device = MagicMock()
        del device.id  # Remove id attribute
        lw = LazyWeight(
            source=mock_tensor,
            device=device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        fingerprint = lw._get_fingerprint()
        assert "device_single" in fingerprint

    def test_fingerprint_changes_with_different_dtype(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint changes with different dtype."""
        lw1 = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            dtype=ttnn.bfloat16,
        )
        lw2 = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            dtype=ttnn.bfloat4_b,
        )
        assert lw1._get_fingerprint() != lw2._get_fingerprint()


# ============================================================================
# Test _get_mesh_mapper_fingerprint Method
# ============================================================================


class TestGetMeshMapperFingerprint:
    """Tests for LazyWeight._get_mesh_mapper_fingerprint() method."""

    def test_fingerprint_replicated_when_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that fingerprint is 'replicated' when mesh_mapper_config is None."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper_config=None,
        )
        fingerprint = lw._get_mesh_mapper_fingerprint()
        assert fingerprint == "replicated"

    def test_fingerprint_with_mesh_mapper_config(
        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
    ):
        """Test that fingerprint includes mapper hash when config is set."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper_config=mock_mesh_mapper_config,
        )
        fingerprint = lw._get_mesh_mapper_fingerprint()
        assert fingerprint.startswith("mapper_")
        # Hash should be 12 characters
        assert len(fingerprint) == len("mapper_") + 12

    def test_different_configs_produce_different_fingerprints(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that different mesh_mapper_configs produce different fingerprints."""
        config1 = ttnn.MeshMapperConfig(
            placements=[ttnn.PlacementShard(-1)],
            mesh_shape_override=ttnn.MeshShape([8]),
        )
        config2 = ttnn.MeshMapperConfig(
            placements=[ttnn.PlacementShard(-2)],
            mesh_shape_override=ttnn.MeshShape([8]),
        )
        lw1 = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper_config=config1,
        )
        lw2 = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper_config=config2,
        )
        assert lw1._get_mesh_mapper_fingerprint() != lw2._get_mesh_mapper_fingerprint()


# ============================================================================
# Test _get_cache_fill_path Method
# ============================================================================


class TestGetCacheFillPath:
    """Tests for LazyWeight._get_cache_fill_path() method."""

    def test_returns_none_when_cache_dir_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that _get_cache_fill_path returns None when cache_dir is None."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        result = lw._get_cache_fill_path(cache_dir=None, weight_name="test")
        assert result is None

    def test_returns_none_when_weight_name_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that _get_cache_fill_path returns None when weight_name is None."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        result = lw._get_cache_fill_path(cache_dir=Path("/tmp/cache"), weight_name=None)
        assert result is None

    def test_returns_path_with_fingerprint(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that _get_cache_fill_path returns correct path with fingerprint."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        cache_dir = Path("/tmp/cache")
        result = lw._get_cache_fill_path(cache_dir=cache_dir, weight_name="my_weight")
        assert result is not None
        assert result.parent == cache_dir
        assert result.name.startswith("my_weight_")
        assert result.suffix == ".tensorbin"


# ============================================================================
# Test padded_shape Property
# ============================================================================


class TestPaddedShape:
    """Tests for LazyWeight.padded_shape property."""

    def test_padded_shape_raises_when_not_resolved(self, mock_tensor, mock_mesh_device):
        """Test that padded_shape raises assertion when not resolved."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            # Missing memory_config and layout
        )
        with pytest.raises(AssertionError, match="LazyWeight must be resolved"):
            _ = lw.padded_shape

    def test_padded_shape_raises_when_device_none(self, mock_tensor, mock_memory_config, mock_mesh_mapper_config):
        """Test that padded_shape raises ValueError when device is None."""
        lw = LazyWeight(
            source=mock_tensor,
            device=None,
            memory_config=mock_memory_config,
            mesh_mapper_config=mock_mesh_mapper_config,
            layout=ttnn.TILE_LAYOUT,
        )
        # First it will fail is_resolved check
        with pytest.raises(AssertionError, match="LazyWeight must be resolved"):
            _ = lw.padded_shape

    def test_padded_shape_no_padding_when_replicated(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that padded_shape returns original shape when replicated (mesh_mapper_config is None)."""
        # Need to set mesh_mapper_config to something to pass is_resolved - use a mock
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper_config=None,
        )
        # Override is_resolved to return True for this test
        with patch.object(lw, "is_resolved", return_value=True):
            result = lw.padded_shape
            assert result == tuple(mock_tensor.shape)

    def test_padded_shape_no_padding_single_device(self, mock_tensor, mock_memory_config, mock_mesh_mapper_config):
        """Test that padded_shape returns original shape for single device."""
        single_device = MagicMock(spec=ttnn.MeshDevice)
        single_device.get_num_devices.return_value = 1

        lw = LazyWeight(
            source=mock_tensor,
            device=single_device,
            memory_config=mock_memory_config,
            mesh_mapper_config=mock_mesh_mapper_config,
            layout=ttnn.TILE_LAYOUT,
        )
        result = lw.padded_shape
        assert result == tuple(mock_tensor.shape)

    def test_padded_shape_pads_for_sharding(self, mock_memory_config):
        """Test that padded_shape correctly pads for sharding across multiple devices."""
        # Create tensor with shape that needs padding: dim -1 is 100
        # With 8 devices, shard_dim = 100/8 = 12.5, needs to be tile aligned (32)
        # padded_shard = ceil(12.5/32)*32 = 32
        # padded_hidden = 32 * 8 = 256
        tensor = torch.randn(1, 32, 100)
        device = MagicMock(spec=ttnn.MeshDevice)
        device.get_num_devices.return_value = 8

        config = ttnn.MeshMapperConfig(
            placements=[ttnn.PlacementShard(-1)],
            mesh_shape_override=ttnn.MeshShape([8]),
        )

        lw = LazyWeight(
            source=tensor,
            device=device,
            memory_config=mock_memory_config,
            mesh_mapper_config=config,
            layout=ttnn.TILE_LAYOUT,
        )
        result = lw.padded_shape
        # Original: (1, 32, 100), padded dim -1 should be 256 (32*8 where 32 is tile-aligned shard)
        assert result == (1, 32, 256)

    def test_padded_shape_raises_valueerror_when_device_none_but_resolved(
        self, mock_tensor, mock_memory_config, mock_mesh_mapper_config
    ):
        """Test that padded_shape raises ValueError when device is None but is_resolved mocked True."""
        lw = LazyWeight(
            source=mock_tensor,
            device=None,
            memory_config=mock_memory_config,
            mesh_mapper_config=mock_mesh_mapper_config,
            layout=ttnn.TILE_LAYOUT,
        )
        # Override is_resolved to return True to reach the ValueError branch
        with patch.object(lw, "is_resolved", return_value=True):
            with pytest.raises(ValueError, match="device must be set to compute padded_shape"):
                _ = lw.padded_shape


# ============================================================================
# Test get_device_weight Method
# ============================================================================


class TestGetDeviceWeight:
    """Tests for LazyWeight.get_device_weight() method."""

    def test_raises_when_device_none(self, mock_tensor, mock_memory_config):
        """Test that get_device_weight raises when device is None."""
        lw = LazyWeight(
            source=mock_tensor,
            device=None,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        with pytest.raises(ValueError, match="device must be provided"):
            lw.get_device_weight()

    def test_raises_when_layout_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that get_device_weight raises when layout is None."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=None,
        )
        with pytest.raises(ValueError, match="layout must be provided"):
            lw.get_device_weight()

    def test_raises_when_memory_config_none(self, mock_tensor, mock_mesh_device):
        """Test that get_device_weight raises when memory_config is None."""
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=None,
            layout=ttnn.TILE_LAYOUT,
        )
        with pytest.raises(ValueError, match="memory_config must be provided"):
            lw.get_device_weight()

    def test_returns_cached_value_on_second_call(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that get_device_weight returns cached _value on second call."""
        mock_ttnn_tensor = MagicMock()
        lw = LazyWeight(
            source=mock_tensor,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )
        # Manually set _value to simulate previous call
        lw._value = mock_ttnn_tensor

        result = lw.get_device_weight()
        assert result is mock_ttnn_tensor

    def test_get_device_weight_replicated_flow(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test get_device_weight with replicated tensor (no sharding) - mocked."""
        mock_ttnn_tensor = MagicMock()
        mock_mapper = MagicMock()

        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
            mock_ttnn.replicate_tensor_to_mesh_mapper.return_value = mock_mapper
            mock_ttnn.from_torch.return_value = mock_ttnn_tensor

            lw = LazyWeight(
                source=mock_tensor,
                device=mock_mesh_device,
                memory_config=mock_memory_config,
                layout=ttnn.TILE_LAYOUT,
                mesh_mapper_config=None,  # Replicated
            )

            result = lw.get_device_weight()

            # Verify replicate_tensor_to_mesh_mapper was called
            mock_ttnn.replicate_tensor_to_mesh_mapper.assert_called_once_with(mock_mesh_device)
            assert result is mock_ttnn_tensor

    def test_get_device_weight_sharded_flow(
        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
    ):
        """Test get_device_weight with sharded tensor - mocked."""
        mock_ttnn_tensor = MagicMock()
        mock_mapper = MagicMock()

        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
            mock_ttnn.create_mesh_mapper.return_value = mock_mapper
            mock_ttnn.from_torch.return_value = mock_ttnn_tensor

            lw = LazyWeight(
                source=mock_tensor,
                device=mock_mesh_device,
                memory_config=mock_memory_config,
                layout=ttnn.TILE_LAYOUT,
                mesh_mapper_config=mock_mesh_mapper_config,
            )

            result = lw.get_device_weight()

            # Verify create_mesh_mapper was called with the config
            mock_ttnn.create_mesh_mapper.assert_called_once_with(mock_mesh_device, mock_mesh_mapper_config)
            assert result is mock_ttnn_tensor

    def test_get_device_weight_with_cache_dir_weight_name(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test get_device_weight extracts cache_dir and weight_name from tuple."""
        mock_ttnn_tensor = MagicMock()
        mock_mapper = MagicMock()

        with tempfile.TemporaryDirectory() as tmpdir:
            cache_dir = Path(tmpdir)

            with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
                mock_ttnn.replicate_tensor_to_mesh_mapper.return_value = mock_mapper
                mock_ttnn.from_torch.return_value = mock_ttnn_tensor
                mock_ttnn.StorageType.HOST = ttnn.StorageType.HOST
                mock_ttnn_tensor.storage_type.return_value = ttnn.StorageType.HOST

                lw = LazyWeight(
                    source=mock_tensor,
                    cache_dir_weight_name=(cache_dir, "my_weight"),
                    device=mock_mesh_device,
                    memory_config=mock_memory_config,
                    layout=ttnn.TILE_LAYOUT,
                    mesh_mapper_config=None,
                )

                result = lw.get_device_weight()
                assert result is not None

    def test_get_device_weight_cache_hit(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test get_device_weight loads from cache when file exists."""
        mock_ttnn_tensor = MagicMock()

        with tempfile.TemporaryDirectory() as tmpdir:
            cache_dir = Path(tmpdir)

            # Create LazyWeight and compute cache path WITHOUT mocking
            lw = LazyWeight(
                source=mock_tensor,
                cache_dir_weight_name=(cache_dir, "cached_weight"),
                device=mock_mesh_device,
                memory_config=mock_memory_config,
                layout=ttnn.TILE_LAYOUT,
                mesh_mapper_config=None,
            )

            # Get expected cache path and create fake cache file
            cache_path = lw._get_cache_fill_path(cache_dir, "cached_weight")
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            cache_path.touch()

            # Now patch ttnn.load_tensor for the get_device_weight call
            with patch("models.common.modules.lazy_weight.ttnn.load_tensor") as mock_load:
                mock_load.return_value = mock_ttnn_tensor

                result = lw.get_device_weight()

                # Verify load_tensor was called (cache hit)
                mock_load.assert_called_once()
                assert result is mock_ttnn_tensor


# ============================================================================
# Test _auto_pad_for_sharded_tiles Function
# ============================================================================


class TestAutoPadForSharding:
    """Tests for _auto_pad_for_sharded_tiles() function."""

    def test_no_padding_when_shapes_match(self, mock_tensor):
        """Test that no padding is applied when shapes match."""
        result = _auto_pad_for_sharded_tiles(mock_tensor, tuple(mock_tensor.shape))
        assert result is mock_tensor

    def test_padding_applied_when_shapes_differ(self):
        """Test that padding is applied when shapes differ."""
        tensor = torch.randn(1, 32, 100)
        padded_shape = (1, 32, 128)

        result = _auto_pad_for_sharded_tiles(tensor, padded_shape)
        assert result.shape == padded_shape
        # Original data should be preserved
        assert torch.equal(result[:, :, :100], tensor)

    def test_padding_with_custom_pad_value(self):
        """Test padding with custom pad value."""
        tensor = torch.zeros(1, 32, 100)
        padded_shape = (1, 32, 128)

        result = _auto_pad_for_sharded_tiles(tensor, padded_shape, pad_value=1.0)
        assert result.shape == padded_shape
        # Padded region should have pad_value
        assert torch.allclose(result[:, :, 100:], torch.ones(1, 32, 28))


# ============================================================================
# Test _from_torch_and_dump Function
# ============================================================================


class TestFromTorchAndDump:
    """Tests for _from_torch_and_dump() function."""

    def test_no_cache_when_cache_file_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that no caching occurs when cache_file_name is None."""
        mock_ttnn_tensor = MagicMock()
        mock_mapper = MagicMock()

        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
            mock_ttnn.from_torch.return_value = mock_ttnn_tensor

            result = _from_torch_and_dump(
                tensor=mock_tensor,
                device=mock_mesh_device,
                dtype=ttnn.bfloat16,
                layout=ttnn.TILE_LAYOUT,
                memory_config=mock_memory_config,
                mesh_mapper=mock_mapper,
                is_replicated=False,
                pad_value=0.0,
                cache_file_name=None,
            )

            mock_ttnn.from_torch.assert_called_once()
            assert result is mock_ttnn_tensor

    def test_cache_dump_when_cache_file_provided(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that tensor is cached when cache_file_name is provided."""
        mock_ttnn_tensor = MagicMock()
        mock_ttnn_tensor.storage_type.return_value = ttnn.StorageType.HOST

        with tempfile.TemporaryDirectory() as tmpdir:
            cache_file = Path(tmpdir) / "test_cache.tensorbin"

            with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
                mock_ttnn.from_torch.return_value = mock_ttnn_tensor
                mock_ttnn.StorageType.HOST = ttnn.StorageType.HOST

                _from_torch_and_dump(
                    tensor=mock_tensor,
                    device=mock_mesh_device,
                    dtype=ttnn.bfloat16,
                    layout=ttnn.TILE_LAYOUT,
                    memory_config=mock_memory_config,
                    mesh_mapper=None,
                    is_replicated=True,
                    pad_value=0.0,
                    cache_file_name=cache_file,
                )

                # Verify from_torch was called with device=None for caching
                call_kwargs = mock_ttnn.from_torch.call_args.kwargs
                assert call_kwargs.get("device") is None

                # Verify dump_tensor_flatbuffer was called
                mock_ttnn._ttnn.tensor.dump_tensor_flatbuffer.assert_called_once()

    def test_replicated_uses_none_mesh_mapper(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that replicated tensors use None mesh_mapper in from_torch."""
        mock_ttnn_tensor = MagicMock()
        mock_mapper = MagicMock()

        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
            mock_ttnn.from_torch.return_value = mock_ttnn_tensor

            _from_torch_and_dump(
                tensor=mock_tensor,
                device=mock_mesh_device,
                dtype=ttnn.bfloat16,
                layout=ttnn.TILE_LAYOUT,
                memory_config=mock_memory_config,
                mesh_mapper=mock_mapper,
                is_replicated=True,
                pad_value=0.0,
                cache_file_name=None,
            )

            call_kwargs = mock_ttnn.from_torch.call_args.kwargs
            assert call_kwargs.get("mesh_mapper") is None

    def test_sharded_uses_mesh_mapper(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that sharded tensors use the provided mesh_mapper."""
        mock_ttnn_tensor = MagicMock()
        mock_mapper = MagicMock()

        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
            mock_ttnn.from_torch.return_value = mock_ttnn_tensor

            _from_torch_and_dump(
                tensor=mock_tensor,
                device=mock_mesh_device,
                dtype=ttnn.bfloat16,
                layout=ttnn.TILE_LAYOUT,
                memory_config=mock_memory_config,
                mesh_mapper=mock_mapper,
                is_replicated=False,
                pad_value=0.0,
                cache_file_name=None,
            )

            call_kwargs = mock_ttnn.from_torch.call_args.kwargs
            assert call_kwargs.get("mesh_mapper") is mock_mapper


# ============================================================================
# Test resolve_lazy_weight Function
# ============================================================================


class TestResolveLazyWeight:
    """Tests for resolve_lazy_weight() function."""

    def test_resolves_none_fields(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that resolve_lazy_weight fills in None fields."""
        lw = LazyWeight(source=mock_tensor)

        resolved = resolve_lazy_weight(
            lw,
            device=mock_mesh_device,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )

        assert resolved.device is mock_mesh_device
        assert resolved.memory_config is mock_memory_config
        assert resolved.layout == ttnn.TILE_LAYOUT

    def test_does_not_override_existing_fields(self, mock_tensor, mock_mesh_device, mock_memory_config):
        """Test that resolve_lazy_weight does not override non-None fields."""
        original_device = MagicMock(spec=ttnn.MeshDevice)
        lw = LazyWeight(
            source=mock_tensor,
            device=original_device,
            dtype=ttnn.bfloat4_b,
        )

        resolved = resolve_lazy_weight(
            lw,
            device=mock_mesh_device,
            dtype=ttnn.bfloat16,
            memory_config=mock_memory_config,
            layout=ttnn.TILE_LAYOUT,
        )

        # Original non-None fields should be preserved
        assert resolved.device is original_device
        assert resolved.dtype == ttnn.bfloat4_b
        # None fields should be filled
        assert resolved.memory_config is mock_memory_config
        assert resolved.layout == ttnn.TILE_LAYOUT

    def test_returns_new_instance(self, mock_tensor, mock_mesh_device):
        """Test that resolve_lazy_weight returns a new LazyWeight instance."""
        lw = LazyWeight(source=mock_tensor)
        resolved = resolve_lazy_weight(lw, device=mock_mesh_device)

        assert resolved is not lw
        assert resolved.device is mock_mesh_device
        assert lw.device is None  # Original unchanged

    def test_with_no_overrides(self, mock_tensor):
        """Test resolve_lazy_weight with no kwargs returns equivalent instance."""
        lw = LazyWeight(source=mock_tensor, dtype=ttnn.bfloat16)
        resolved = resolve_lazy_weight(lw)

        assert resolved is not lw
        assert resolved.source is lw.source
        assert resolved.dtype == lw.dtype


# ============================================================================
# Integration Tests (require device - use parametrized fixture)
# ============================================================================


@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [
        (1, 1),
    ],
    ids=[
        "1x1",
    ],
    indirect=True,
)
class TestLazyWeightIntegration:
    """Integration tests for LazyWeight that require actual TTNN device."""

    def test_get_device_weight_replicated(self, ttnn_mesh_device: ttnn.MeshDevice):
        """Test get_device_weight with replicated tensor (no mesh_mapper_config)."""
        tensor = torch.randn(1, 1, 32, 64, dtype=torch.bfloat16)

        lw = LazyWeight(
            source=tensor,
            device=ttnn_mesh_device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper_config=None,  # Replicated
        )

        result = lw.get_device_weight()
        assert result is not None
        # Subsequent calls should return cached value
        result2 = lw.get_device_weight()
        assert result2 is result

    def test_get_device_weight_with_cache(self, ttnn_mesh_device: ttnn.MeshDevice):
        """Test get_device_weight with caching enabled."""
        tensor = torch.randn(1, 1, 32, 64, dtype=torch.bfloat16)

        with tempfile.TemporaryDirectory() as tmpdir:
            cache_dir = Path(tmpdir)

            lw = LazyWeight(
                source=tensor,
                cache_dir_weight_name=(cache_dir, "test_weight"),
                device=ttnn_mesh_device,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                layout=ttnn.TILE_LAYOUT,
                mesh_mapper_config=None,
            )

            # First call should create cache
            result1 = lw.get_device_weight()
            assert result1 is not None

            # Verify cache file was created
            cache_files = list(cache_dir.glob("*.tensorbin"))
            assert len(cache_files) == 1
            assert cache_files[0].name.startswith("test_weight_")

    def test_get_device_weight_cache_hit(self, ttnn_mesh_device: ttnn.MeshDevice):
        """Test get_device_weight loads from cache on second instance."""
        tensor = torch.randn(1, 1, 32, 64, dtype=torch.bfloat16)

        with tempfile.TemporaryDirectory() as tmpdir:
            cache_dir = Path(tmpdir)

            # First LazyWeight creates cache
            lw1 = LazyWeight(
                source=tensor,
                cache_dir_weight_name=(cache_dir, "test_weight"),
                device=ttnn_mesh_device,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                layout=ttnn.TILE_LAYOUT,
                mesh_mapper_config=None,
            )
            lw1.get_device_weight()

            # Second LazyWeight should load from cache
            lw2 = LazyWeight(
                source=tensor,
                cache_dir_weight_name=(cache_dir, "test_weight"),
                device=ttnn_mesh_device,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                layout=ttnn.TILE_LAYOUT,
                mesh_mapper_config=None,
            )
            result2 = lw2.get_device_weight()

            assert result2 is not None


if __name__ == "__main__":
    # Run unit tests that don't require device
    pytest.main([__file__, "-v", "-k", "not Integration"])