code/models/tt_dit/tests/unit/test_tensor_utils.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /tests /unit /test_tensor_utils.py

Add files using upload-large-folder tool

33d23e1 verified 11 days ago

4.28 kB

# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC

# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import pytest
import torch

import ttnn

from ...utils import tensor
from ...utils.tracing import Tracer


@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_full(mesh_device: ttnn.MeshDevice) -> None:
    shape = (32, 32)
    dtype = ttnn.bfloat16
    value = 10

    full = Tracer(
        lambda: tensor.full(shape, value, dtype=dtype, device=mesh_device),
        device=mesh_device,
    )

    result = full()
    ref = torch.full(shape, value, dtype=torch.bfloat16)

    assert result.dtype == dtype
    assert tuple(result.shape) == tuple(ref.shape)

    result_torch = tensor.to_torch(result)
    assert torch.allclose(result_torch, ref, atol=0, rtol=0)


@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_arange(mesh_device: ttnn.MeshDevice) -> None:
    dtype = ttnn.bfloat16
    start, end, step = 10, 20, 2

    arange = Tracer(
        lambda: tensor.arange(start, end, step, dtype=dtype, device=mesh_device),
        device=mesh_device,
    )

    result = arange()
    ref = torch.arange(start, end, step, dtype=torch.bfloat16)

    assert result.dtype == dtype
    assert tuple(result.shape) == tuple(ref.shape)

    result_torch = tensor.to_torch(result)
    assert torch.allclose(result_torch, ref, atol=0, rtol=0)


@pytest.mark.parametrize(
    ("shape", "dim", "front", "back"),
    [
        # On device tile padding does not support front padding.
        # rank <= 4: direct ttnn.pad path
        ((2, 32, 64), 0, 0, 2),
        ((2, 32, 64), 1, 0, 3),
        # rank > 4, dim >= rank - 3: direct ttnn.pad path (last 3 dims)
        ((2, 2, 2, 32, 64), 3, 0, 2),
        ((2, 2, 2, 32, 64), -1, 0, 4),
        # rank > 4, dim < rank - 3: reshape workaround path (early dims)
        ((2, 2, 2, 32, 64), 0, 0, 2),
        ((2, 2, 2, 32, 64), 1, 0, 3),
    ],
)
@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_pad_single(mesh_device: ttnn.MeshDevice, shape: tuple, dim: int, front: int, back: int) -> None:
    dtype = ttnn.bfloat16

    x_torch = torch.randn(shape, dtype=torch.bfloat16)
    x_tt = tensor.from_torch(x_torch, device=mesh_device, dtype=dtype)

    pad = Tracer(
        lambda: tensor.pad_single(x_tt, dim=dim, front=front, back=back),
        device=mesh_device,
    )

    result = pad()

    pad_per_dim = [(front, back) if i == (dim % len(shape)) else (0, 0) for i in reversed(range(len(shape)))]
    ref = torch.nn.functional.pad(x_torch, [v for pair in pad_per_dim for v in pair])

    assert result.dtype == dtype
    assert tuple(result.shape) == tuple(ref.shape)

    result_torch = tensor.to_torch(result)
    assert torch.allclose(result_torch, ref, atol=0, rtol=0)


@pytest.mark.parametrize("diagonal", [0, 1, -1])
@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_tril(mesh_device: ttnn.MeshDevice, diagonal: int) -> None:
    shape = (1, 1, 64, 64)
    dtype = ttnn.bfloat16

    x_torch = torch.ones(shape, dtype=torch.bfloat16)
    x_tt = tensor.from_torch(x_torch, device=mesh_device, dtype=dtype)

    tril = Tracer(
        lambda: tensor.tril(x_tt, diagonal=diagonal),
        device=mesh_device,
    )

    result = tril()
    ref = torch.tril(x_torch, diagonal=diagonal)

    assert result.dtype == dtype
    assert tuple(result.shape) == tuple(ref.shape)

    result_torch = tensor.to_torch(result)
    assert torch.allclose(result_torch, ref, atol=0, rtol=0)


@pytest.mark.parametrize("diagonal", [0, 1, -1])
@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_triu(mesh_device: ttnn.MeshDevice, diagonal: int) -> None:
    shape = (1, 1, 64, 64)
    dtype = ttnn.bfloat16

    x_torch = torch.ones(shape, dtype=torch.bfloat16)
    x_tt = tensor.from_torch(x_torch, device=mesh_device, dtype=dtype)

    triu = Tracer(
        lambda: tensor.triu(x_tt, diagonal=diagonal),
        device=mesh_device,
    )

    result = triu()
    ref = torch.triu(x_torch, diagonal=diagonal)

    assert result.dtype == dtype
    assert tuple(result.shape) == tuple(ref.shape)

    result_torch = tensor.to_torch(result)
    assert torch.allclose(result_torch, ref, atol=0, rtol=0)

	# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC

	# SPDX-License-Identifier: Apache-2.0

	from __future__ import annotations

	import pytest
	import torch

	import ttnn

	from ...utils import tensor
	from ...utils.tracing import Tracer


	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_full(mesh_device: ttnn.MeshDevice) -> None:
	shape = (32, 32)
	dtype = ttnn.bfloat16
	value = 10

	full = Tracer(
	lambda: tensor.full(shape, value, dtype=dtype, device=mesh_device),
	device=mesh_device,
	)

	result = full()
	ref = torch.full(shape, value, dtype=torch.bfloat16)

	assert result.dtype == dtype
	assert tuple(result.shape) == tuple(ref.shape)

	result_torch = tensor.to_torch(result)
	assert torch.allclose(result_torch, ref, atol=0, rtol=0)


	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_arange(mesh_device: ttnn.MeshDevice) -> None:
	dtype = ttnn.bfloat16
	start, end, step = 10, 20, 2

	arange = Tracer(
	lambda: tensor.arange(start, end, step, dtype=dtype, device=mesh_device),
	device=mesh_device,
	)

	result = arange()
	ref = torch.arange(start, end, step, dtype=torch.bfloat16)

	assert result.dtype == dtype
	assert tuple(result.shape) == tuple(ref.shape)

	result_torch = tensor.to_torch(result)
	assert torch.allclose(result_torch, ref, atol=0, rtol=0)


	@pytest.mark.parametrize(
	("shape", "dim", "front", "back"),
	[
	# On device tile padding does not support front padding.
	# rank <= 4: direct ttnn.pad path
	((2, 32, 64), 0, 0, 2),
	((2, 32, 64), 1, 0, 3),
	# rank > 4, dim >= rank - 3: direct ttnn.pad path (last 3 dims)
	((2, 2, 2, 32, 64), 3, 0, 2),
	((2, 2, 2, 32, 64), -1, 0, 4),
	# rank > 4, dim < rank - 3: reshape workaround path (early dims)
	((2, 2, 2, 32, 64), 0, 0, 2),
	((2, 2, 2, 32, 64), 1, 0, 3),
	],
	)
	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_pad_single(mesh_device: ttnn.MeshDevice, shape: tuple, dim: int, front: int, back: int) -> None:
	dtype = ttnn.bfloat16

	x_torch = torch.randn(shape, dtype=torch.bfloat16)
	x_tt = tensor.from_torch(x_torch, device=mesh_device, dtype=dtype)

	pad = Tracer(
	lambda: tensor.pad_single(x_tt, dim=dim, front=front, back=back),
	device=mesh_device,
	)

	result = pad()

	pad_per_dim = [(front, back) if i == (dim % len(shape)) else (0, 0) for i in reversed(range(len(shape)))]
	ref = torch.nn.functional.pad(x_torch, [v for pair in pad_per_dim for v in pair])

	assert result.dtype == dtype
	assert tuple(result.shape) == tuple(ref.shape)

	result_torch = tensor.to_torch(result)
	assert torch.allclose(result_torch, ref, atol=0, rtol=0)


	@pytest.mark.parametrize("diagonal", [0, 1, -1])
	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_tril(mesh_device: ttnn.MeshDevice, diagonal: int) -> None:
	shape = (1, 1, 64, 64)
	dtype = ttnn.bfloat16

	x_torch = torch.ones(shape, dtype=torch.bfloat16)
	x_tt = tensor.from_torch(x_torch, device=mesh_device, dtype=dtype)

	tril = Tracer(
	lambda: tensor.tril(x_tt, diagonal=diagonal),
	device=mesh_device,
	)

	result = tril()
	ref = torch.tril(x_torch, diagonal=diagonal)

	assert result.dtype == dtype
	assert tuple(result.shape) == tuple(ref.shape)

	result_torch = tensor.to_torch(result)
	assert torch.allclose(result_torch, ref, atol=0, rtol=0)


	@pytest.mark.parametrize("diagonal", [0, 1, -1])
	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_triu(mesh_device: ttnn.MeshDevice, diagonal: int) -> None:
	shape = (1, 1, 64, 64)
	dtype = ttnn.bfloat16

	x_torch = torch.ones(shape, dtype=torch.bfloat16)
	x_tt = tensor.from_torch(x_torch, device=mesh_device, dtype=dtype)

	triu = Tracer(
	lambda: tensor.triu(x_tt, diagonal=diagonal),
	device=mesh_device,
	)

	result = triu()
	ref = torch.triu(x_torch, diagonal=diagonal)

	assert result.dtype == dtype
	assert tuple(result.shape) == tuple(ref.shape)

	result_torch = tensor.to_torch(result)
	assert torch.allclose(result_torch, ref, atol=0, rtol=0)