xhs-note-crawling/python/py/Lib/site-packages/sympy/printing/tests/test_torch.py

import random
import math

from sympy import symbols, Derivative
from sympy.printing.pytorch import torch_code
from sympy import (eye, MatrixSymbol, Matrix)
from sympy.tensor.array import NDimArray
from sympy.tensor.array.expressions.array_expressions import (
    ArrayTensorProduct, ArrayAdd,
    PermuteDims, ArrayDiagonal, _CodegenArrayAbstract)
from sympy.utilities.lambdify import lambdify
from sympy.core.relational import Eq, Ne, Ge, Gt, Le, Lt
from sympy.functions import \
    Abs, ceiling, exp, floor, sign, sin, asin, cos, \
    acos, tan, atan, atan2, cosh, acosh, sinh, asinh, tanh, atanh, \
    re, im, arg, erf, loggamma, sqrt
from sympy.testing.pytest import skip
from sympy.external import import_module
from sympy.matrices.expressions import \
    Determinant, HadamardProduct, Inverse, Trace
from sympy.matrices import randMatrix
from sympy.matrices import Identity, ZeroMatrix, OneMatrix
from sympy import conjugate, I
from sympy import Heaviside, gamma, polygamma



torch = import_module("torch")

M = MatrixSymbol("M", 3, 3)
N = MatrixSymbol("N", 3, 3)
P = MatrixSymbol("P", 3, 3)
Q = MatrixSymbol("Q", 3, 3)

x, y, z, t = symbols("x y z t")

if torch is not None:
    llo = [list(range(i, i + 3)) for i in range(0, 9, 3)]
    m3x3 = torch.tensor(llo, dtype=torch.float64)
    m3x3sympy = Matrix(llo)


def _compare_torch_matrix(variables, expr):
    f = lambdify(variables, expr, 'torch')

    random_matrices = [randMatrix(i.shape[0], i.shape[1]) for i in variables]
    random_variables = [torch.tensor(i.tolist(), dtype=torch.float64) for i in random_matrices]
    r = f(*random_variables)
    e = expr.subs(dict(zip(variables, random_matrices))).doit()

    if isinstance(e, _CodegenArrayAbstract):
        e = e.doit()

    if hasattr(e, 'is_number') and e.is_number:
        if isinstance(r, torch.Tensor) and r.dim() == 0:
            r = r.item()
            e = float(e)
            assert abs(r - e) < 1e-6
            return

    if e.is_Matrix or isinstance(e, NDimArray):
        e = torch.tensor(e.tolist(), dtype=torch.float64)
        assert torch.allclose(r, e, atol=1e-6)
    else:
        raise TypeError(f"Cannot compare {type(r)} with {type(e)}")


def _compare_torch_scalar(variables, expr, rng=lambda: random.uniform(-5, 5)):
    f = lambdify(variables, expr, 'torch')
    rvs = [rng() for v in variables]
    t_rvs = [torch.tensor(i, dtype=torch.float64) for i in rvs]
    r = f(*t_rvs)
    if isinstance(r, torch.Tensor):
        r = r.item()
    e = expr.subs(dict(zip(variables, rvs))).doit()
    assert abs(r - e) < 1e-6


def _compare_torch_relational(variables, expr, rng=lambda: random.randint(0, 10)):
    f = lambdify(variables, expr, 'torch')
    rvs = [rng() for v in variables]
    t_rvs = [torch.tensor(i, dtype=torch.float64) for i in rvs]
    r = f(*t_rvs)
    e = bool(expr.subs(dict(zip(variables, rvs))).doit())
    assert r.item() == e


def test_torch_math():
    if not torch:
        skip("PyTorch not installed")

    expr = Abs(x)
    assert torch_code(expr) == "torch.abs(x)"
    f = lambdify(x, expr, 'torch')
    ma = torch.tensor([[-1, 2, -3, -4]], dtype=torch.float64)
    y_abs = f(ma)
    c = torch.abs(ma)
    assert torch.all(y_abs == c)

    expr = sign(x)
    assert torch_code(expr) == "torch.sign(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-10, 10))

    expr = ceiling(x)
    assert torch_code(expr) == "torch.ceil(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.random())

    expr = floor(x)
    assert torch_code(expr) == "torch.floor(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.random())

    expr = exp(x)
    assert torch_code(expr) == "torch.exp(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-2, 2))

    expr = sqrt(x)
    assert torch_code(expr) == "torch.sqrt(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.random())

    expr = x ** 4
    assert torch_code(expr) == "torch.pow(x, 4)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.random())

    expr = cos(x)
    assert torch_code(expr) == "torch.cos(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.random())

    expr = acos(x)
    assert torch_code(expr) == "torch.acos(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-0.99, 0.99))

    expr = sin(x)
    assert torch_code(expr) == "torch.sin(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.random())

    expr = asin(x)
    assert torch_code(expr) == "torch.asin(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-0.99, 0.99))

    expr = tan(x)
    assert torch_code(expr) == "torch.tan(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-1.5, 1.5))

    expr = atan(x)
    assert torch_code(expr) == "torch.atan(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-5, 5))

    expr = atan2(y, x)
    assert torch_code(expr) == "torch.atan2(y, x)"
    _compare_torch_scalar((y, x), expr, rng=lambda: random.uniform(-5, 5))

    expr = cosh(x)
    assert torch_code(expr) == "torch.cosh(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-2, 2))

    expr = acosh(x)
    assert torch_code(expr) == "torch.acosh(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(1.1, 5))

    expr = sinh(x)
    assert torch_code(expr) == "torch.sinh(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-2, 2))

    expr = asinh(x)
    assert torch_code(expr) == "torch.asinh(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-5, 5))

    expr = tanh(x)
    assert torch_code(expr) == "torch.tanh(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-2, 2))

    expr = atanh(x)
    assert torch_code(expr) == "torch.atanh(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-0.9, 0.9))

    expr = erf(x)
    assert torch_code(expr) == "torch.erf(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(-2, 2))

    expr = loggamma(x)
    assert torch_code(expr) == "torch.lgamma(x)"
    _compare_torch_scalar((x,), expr, rng=lambda: random.uniform(0.5, 5))


def test_torch_complexes():
    assert torch_code(re(x)) == "torch.real(x)"
    assert torch_code(im(x)) == "torch.imag(x)"
    assert torch_code(arg(x)) == "torch.angle(x)"


def test_torch_relational():
    if not torch:
        skip("PyTorch not installed")

    expr = Eq(x, y)
    assert torch_code(expr) == "torch.eq(x, y)"
    _compare_torch_relational((x, y), expr)

    expr = Ne(x, y)
    assert torch_code(expr) == "torch.ne(x, y)"
    _compare_torch_relational((x, y), expr)

    expr = Ge(x, y)
    assert torch_code(expr) == "torch.ge(x, y)"
    _compare_torch_relational((x, y), expr)

    expr = Gt(x, y)
    assert torch_code(expr) == "torch.gt(x, y)"
    _compare_torch_relational((x, y), expr)

    expr = Le(x, y)
    assert torch_code(expr) == "torch.le(x, y)"
    _compare_torch_relational((x, y), expr)

    expr = Lt(x, y)
    assert torch_code(expr) == "torch.lt(x, y)"
    _compare_torch_relational((x, y), expr)


def test_torch_matrix():
    if torch is None:
        skip("PyTorch not installed")

    expr = M
    assert torch_code(expr) == "M"
    f = lambdify((M,), expr, "torch")
    eye_mat = eye(3)
    eye_tensor = torch.tensor(eye_mat.tolist(), dtype=torch.float64)
    assert torch.allclose(f(eye_tensor), eye_tensor)

    expr = M * N
    assert torch_code(expr) == "torch.matmul(M, N)"
    _compare_torch_matrix((M, N), expr)

    expr = M ** 3
    assert torch_code(expr) == "torch.mm(torch.mm(M, M), M)"
    _compare_torch_matrix((M,), expr)

    expr = M * N * P * Q
    assert torch_code(expr) == "torch.matmul(torch.matmul(torch.matmul(M, N), P), Q)"
    _compare_torch_matrix((M, N, P, Q), expr)

    expr = Trace(M)
    assert torch_code(expr) == "torch.trace(M)"
    _compare_torch_matrix((M,), expr)

    expr = Determinant(M)
    assert torch_code(expr) == "torch.det(M)"
    _compare_torch_matrix((M,), expr)

    expr = HadamardProduct(M, N)
    assert torch_code(expr) == "torch.mul(M, N)"
    _compare_torch_matrix((M, N), expr)

    expr = Inverse(M)
    assert torch_code(expr) == "torch.linalg.inv(M)"

    # For inverse, use a matrix that's guaranteed to be invertible
    eye_mat = eye(3)
    eye_tensor = torch.tensor(eye_mat.tolist(), dtype=torch.float64)
    f = lambdify((M,), expr, "torch")
    result = f(eye_tensor)
    expected = torch.linalg.inv(eye_tensor)
    assert torch.allclose(result, expected)


def test_torch_array_operations():
    if not torch:
        skip("PyTorch not installed")

    M = MatrixSymbol("M", 2, 2)
    N = MatrixSymbol("N", 2, 2)
    P = MatrixSymbol("P", 2, 2)
    Q = MatrixSymbol("Q", 2, 2)

    ma = torch.tensor([[1., 2.], [3., 4.]], dtype=torch.float64)
    mb = torch.tensor([[1., -2.], [-1., 3.]], dtype=torch.float64)
    mc = torch.tensor([[2., 0.], [1., 2.]], dtype=torch.float64)
    md = torch.tensor([[1., -1.], [4., 7.]], dtype=torch.float64)

    cg = ArrayTensorProduct(M, N)
    assert torch_code(cg) == 'torch.einsum("ab,cd", M, N)'
    f = lambdify((M, N), cg, 'torch')
    y = f(ma, mb)
    c = torch.einsum("ij,kl", ma, mb)
    assert torch.allclose(y, c)

    cg = ArrayAdd(M, N)
    assert torch_code(cg) == 'torch.add(M, N)'
    f = lambdify((M, N), cg, 'torch')
    y = f(ma, mb)
    c = ma + mb
    assert torch.allclose(y, c)

    cg = ArrayAdd(M, N, P)
    assert torch_code(cg) == 'torch.add(torch.add(M, N), P)'
    f = lambdify((M, N, P), cg, 'torch')
    y = f(ma, mb, mc)
    c = ma + mb + mc
    assert torch.allclose(y, c)

    cg = ArrayAdd(M, N, P, Q)
    assert torch_code(cg) == 'torch.add(torch.add(torch.add(M, N), P), Q)'
    f = lambdify((M, N, P, Q), cg, 'torch')
    y = f(ma, mb, mc, md)
    c = ma + mb + mc + md
    assert torch.allclose(y, c)

    cg = PermuteDims(M, [1, 0])
    assert torch_code(cg) == 'M.permute(1, 0)'
    f = lambdify((M,), cg, 'torch')
    y = f(ma)
    c = ma.T
    assert torch.allclose(y, c)

    cg = PermuteDims(ArrayTensorProduct(M, N), [1, 2, 3, 0])
    assert torch_code(cg) == 'torch.einsum("ab,cd", M, N).permute(1, 2, 3, 0)'
    f = lambdify((M, N), cg, 'torch')
    y = f(ma, mb)
    c = torch.einsum("ab,cd", ma, mb).permute(1, 2, 3, 0)
    assert torch.allclose(y, c)

    cg = ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2))
    assert torch_code(cg) == 'torch.einsum("ab,bc->acb", M, N)'
    f = lambdify((M, N), cg, 'torch')
    y = f(ma, mb)
    c = torch.einsum("ab,bc->acb", ma, mb)
    assert torch.allclose(y, c)


def test_torch_derivative():
    """Test derivative handling."""
    expr = Derivative(sin(x), x)
    assert torch_code(expr) == 'torch.autograd.grad(torch.sin(x), x)[0]'


def test_torch_printing_dtype():
    if not torch:
        skip("PyTorch not installed")

    # matrix printing with default dtype
    expr = Matrix([[x, sin(y)], [exp(z), -t]])
    assert "dtype=torch.float64" in torch_code(expr)

    # explicit dtype
    assert "dtype=torch.float32" in torch_code(expr, dtype="torch.float32")

    # with requires_grad
    result = torch_code(expr, requires_grad=True)
    assert "requires_grad=True" in result
    assert "dtype=torch.float64" in result

    # both
    result = torch_code(expr, requires_grad=True, dtype="torch.float32")
    assert "requires_grad=True" in result
    assert "dtype=torch.float32" in result


def test_requires_grad():
    if not torch:
        skip("PyTorch not installed")

    expr = sin(x) + cos(y)
    f = lambdify([x, y], expr, 'torch')

    # make sure the gradients flow
    x_val = torch.tensor(1.0, requires_grad=True)
    y_val = torch.tensor(2.0, requires_grad=True)
    result = f(x_val, y_val)
    assert result.requires_grad
    result.backward()

    # x_val.grad should be cos(x_val) which is close to cos(1.0)
    assert abs(x_val.grad.item() - float(cos(1.0).evalf())) < 1e-6

    # y_val.grad should be -sin(y_val) which is close to -sin(2.0)
    assert abs(y_val.grad.item() - float(-sin(2.0).evalf())) < 1e-6


def test_torch_multi_variable_derivatives():
    if not torch:
        skip("PyTorch not installed")

    x, y, z = symbols("x y z")

    expr = Derivative(sin(x), x)
    assert torch_code(expr) == "torch.autograd.grad(torch.sin(x), x)[0]"

    expr = Derivative(sin(x), (x, 2))
    assert torch_code(
        expr) == "torch.autograd.grad(torch.autograd.grad(torch.sin(x), x, create_graph=True)[0], x, create_graph=True)[0]"

    expr = Derivative(sin(x * y), x, y)
    result = torch_code(expr)
    expected = "torch.autograd.grad(torch.autograd.grad(torch.sin(x*y), x, create_graph=True)[0], y, create_graph=True)[0]"
    normalized_result = result.replace(" ", "")
    normalized_expected = expected.replace(" ", "")
    assert normalized_result == normalized_expected

    expr = Derivative(sin(x), x, x)
    result = torch_code(expr)
    expected = "torch.autograd.grad(torch.autograd.grad(torch.sin(x), x, create_graph=True)[0], x, create_graph=True)[0]"
    assert result == expected

    expr = Derivative(sin(x * y * z), x, (y, 2), z)
    result = torch_code(expr)
    expected = "torch.autograd.grad(torch.autograd.grad(torch.autograd.grad(torch.autograd.grad(torch.sin(x*y*z), x, create_graph=True)[0], y, create_graph=True)[0], y, create_graph=True)[0], z, create_graph=True)[0]"
    normalized_result = result.replace(" ", "")
    normalized_expected = expected.replace(" ", "")
    assert normalized_result == normalized_expected


def test_torch_derivative_lambdify():
    if not torch:
        skip("PyTorch not installed")

    x = symbols("x")
    y = symbols("y")

    expr = Derivative(x ** 2, x)
    f = lambdify(x, expr, 'torch')
    x_val = torch.tensor(2.0, requires_grad=True)
    result = f(x_val)
    assert torch.isclose(result, torch.tensor(4.0))

    expr = Derivative(sin(x), (x, 2))
    f = lambdify(x, expr, 'torch')
    # Second derivative of sin(x) at x=0 is 0, not -1
    x_val = torch.tensor(0.0, requires_grad=True)
    result = f(x_val)
    assert torch.isclose(result, torch.tensor(0.0), atol=1e-5)

    x_val = torch.tensor(math.pi / 2, requires_grad=True)
    result = f(x_val)
    assert torch.isclose(result, torch.tensor(-1.0), atol=1e-5)

    expr = Derivative(x * y ** 2, x, y)
    f = lambdify((x, y), expr, 'torch')
    x_val = torch.tensor(2.0, requires_grad=True)
    y_val = torch.tensor(3.0, requires_grad=True)
    result = f(x_val, y_val)
    assert torch.isclose(result, torch.tensor(6.0))


def test_torch_special_matrices():
    if not torch:
        skip("PyTorch not installed")

    expr = Identity(3)
    assert torch_code(expr) == "torch.eye(3)"

    n = symbols("n")
    expr = Identity(n)
    assert torch_code(expr) == "torch.eye(n, n)"

    expr = ZeroMatrix(2, 3)
    assert torch_code(expr) == "torch.zeros((2, 3))"

    m, n = symbols("m n")
    expr = ZeroMatrix(m, n)
    assert torch_code(expr) == "torch.zeros((m, n))"

    expr = OneMatrix(2, 3)
    assert torch_code(expr) == "torch.ones((2, 3))"

    expr = OneMatrix(m, n)
    assert torch_code(expr) == "torch.ones((m, n))"


def test_torch_special_matrices_lambdify():
    if not torch:
        skip("PyTorch not installed")

    expr = Identity(3)
    f = lambdify([], expr, 'torch')
    result = f()
    expected = torch.eye(3)
    assert torch.allclose(result, expected)

    expr = ZeroMatrix(2, 3)
    f = lambdify([], expr, 'torch')
    result = f()
    expected = torch.zeros((2, 3))
    assert torch.allclose(result, expected)

    expr = OneMatrix(2, 3)
    f = lambdify([], expr, 'torch')
    result = f()
    expected = torch.ones((2, 3))
    assert torch.allclose(result, expected)


def test_torch_complex_operations():
    if not torch:
        skip("PyTorch not installed")

    expr = conjugate(x)
    assert torch_code(expr) == "torch.conj(x)"

    # SymPy distributes conjugate over addition and applies specific rules for each term
    expr = conjugate(sin(x) + I * cos(y))
    assert torch_code(expr) == "torch.sin(torch.conj(x)) - 1j*torch.cos(torch.conj(y))"

    expr = I
    assert torch_code(expr) == "1j"

    expr = 2 * I + x
    assert torch_code(expr) == "x + 2*1j"

    expr = exp(I * x)
    assert torch_code(expr) == "torch.exp(1j*x)"


def test_torch_special_functions():
    if not torch:
        skip("PyTorch not installed")

    expr = Heaviside(x)
    assert torch_code(expr) == "torch.heaviside(x, 1/2)"

    expr = Heaviside(x, 0)
    assert torch_code(expr) == "torch.heaviside(x, 0)"

    expr = gamma(x)
    assert torch_code(expr) == "torch.special.gamma(x)"

    expr = polygamma(0, x)  # Use polygamma instead of digamma because sympy will default to that anyway
    assert torch_code(expr) == "torch.special.digamma(x)"

    expr = gamma(sin(x))
    assert torch_code(expr) == "torch.special.gamma(torch.sin(x))"