File size: 1,251 Bytes
179df43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | # SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
import pytest
import torch
import ttnn
from tt_sam3.utils import assert_with_pcc
@pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
@pytest.mark.parametrize("batch_size", [1])
def test_tt_patch_embed(device, reset_seeds, batch_size, sam3_vit_backbone):
"""Test PatchEmbed against PyTorch reference."""
from tt_sam3.tt.ttnn_sam3_vitdet import tt_patch_embed, get_patch_embed_params
torch.manual_seed(42)
# Reference
ref_model = sam3_vit_backbone
pixel_values = torch.randn(batch_size, 3, 1008, 1008)
with torch.no_grad():
ref_output = ref_model.patch_embed(pixel_values) # (B, H, W, C) = (1, 72, 72, 1024)
# Our implementation
params = get_patch_embed_params(ref_model)
tt_output = tt_patch_embed(pixel_values, params["weight"], params["bias"], device)
tt_output_torch = ttnn.to_torch(tt_output)
# Compare - shape should match
assert tt_output_torch.shape[-3:] == ref_output.shape[-3:], (
f"Shape mismatch: {tt_output_torch.shape} vs {ref_output.shape}"
)
# PCC check
assert_with_pcc(ref_output.float(), tt_output_torch.float(), 0.999)
|