# Owner(s): ["module: tests"]

# ruff: noqa: F841

import itertools

import math

import operator

import random

import sys

import warnings

from functools import partial

from itertools import chain, product

from numbers import Number

import numpy as np

import torch

import torch.autograd.forward_ad as fwAD

from torch import inf, nan

from torch.testing import make_tensor

from torch.testing._internal.common_device_type import (

deviceCountAtLeast,

dtypes,

dtypesIfCPU,

dtypesIfCUDA,

dtypesIfXPU,

expectedFailureMeta,

instantiate_device_type_tests,

onlyNativeDeviceTypes,

onlyOn,

OpDTypes,

ops,

precisionOverride,

skipIf,

skipMeta,

skipXPU,

)

from torch.testing._internal.common_dtype import (

all_types,

all_types_and,

all_types_and_complex_and,

complex_types,

floating_and_complex_types,

floating_types_and,

get_all_int_dtypes,

get_all_math_dtypes,

integral_types,

integral_types_and,

)

from torch.testing._internal.common_methods_invocations import (

binary_ufuncs,

generate_elementwise_binary_broadcasting_tensors,

generate_elementwise_binary_extremal_value_tensors,

generate_elementwise_binary_large_value_tensors,

generate_elementwise_binary_small_value_tensors,

generate_elementwise_binary_tensors,

generate_elementwise_binary_with_scalar_and_type_promotion_samples,

generate_elementwise_binary_with_scalar_samples,

)

from torch.testing._internal.common_utils import (

gradcheck,

instantiate_parametrized_tests,

iter_indices,

numpy_to_torch_dtype_dict,

parametrize,

run_tests,

set_default_dtype,

skipIfTorchDynamo,

slowTest,

TEST_SCIPY,

TestCase,

torch_to_numpy_dtype_dict,

xfailIfTorchDynamo,

)

if TEST_SCIPY:

import scipy.integrate

import scipy.special

device_type = (

acc.type if (acc := torch.accelerator.current_accelerator(True)) else "cpu"

)

_unsigned_int_types = (torch.uint16, torch.uint32, torch.uint64)

# TODO: update to use opinfos consistently

@instantiate_parametrized_tests

class TestBinaryUfuncs(TestCase):

def _test_cop(self, torchfn, mathfn, dtype):

def reference_implementation(res2):

for i, j in iter_indices(sm1):

idx1d = i * sm1.size(0) + j

res2[i, j] = mathfn(sm1[i, j], sm2[idx1d])

return res2

# contiguous

m1 = torch.randn(10, 10, 10, dtype=dtype)

m2 = torch.randn(10, 10 * 10, dtype=dtype)

sm1 = m1[4]

sm2 = m2[4]

res1 = torchfn(sm1, sm2.view(10, 10))

res2 = reference_implementation(res1.clone())

self.assertEqual(res1, res2)

# non-contiguous

m1 = torch.randn(10, 10, 10, dtype=dtype)

m2 = torch.randn(10 * 10, 10 * 10, dtype=dtype)

sm1 = m1[:, 4]

sm2 = m2[:, 4]

# view as sm1.size()

sm2.set_(

sm2.storage(),

sm2.storage_offset(),

sm1.size(),

(sm2.stride()[0] * 10, sm2.stride()[0]),

)

res1 = torchfn(sm1, sm2)

# reference_implementation assumes 1-d sm2

sm2.set_(

sm2.storage(), sm2.storage_offset(), m2[:, 4].size(), m2[:, 4].stride()

)

res2 = reference_implementation(res1.clone())

self.assertEqual(res1, res2)

@parametrize(

"dtype", list(all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool))

)

def test_sub(self, dtype):

if dtype in integral_types():

# Before Python 3.10, floats were implicitly converted to ints, but with

# DeprecationWarning: an integer is required (got type float).

# Implicit conversion to integers using __int__ is deprecated,

# and may be removed in a future version of Python.

# Since Python 3.10, that attempt gives an error.

m1 = torch.tensor([2, 4], dtype=dtype)

m2 = torch.tensor([1, 2], dtype=dtype)

diff = torch.tensor([1, 2], dtype=dtype)

else:

m1 = torch.tensor([2.34, 4.44], dtype=dtype)

m2 = torch.tensor([1.23, 2.33], dtype=dtype)

diff = torch.tensor([1.11, 2.11], dtype=dtype)

if dtype == torch.bool:

self.assertRaises(RuntimeError, lambda: m1 - m2)

elif dtype == torch.bfloat16 or dtype == torch.half:

# bfloat16 has a lower precision so we have to have a separate check for it

self.assertEqual(m1 - m2, diff, atol=0.01, rtol=0)

else:

self.assertEqual(m1 - m2, diff)

@parametrize("dtype", [torch.float])

def test_csub(self, dtype):

# with a tensor

a = torch.randn(100, 90, dtype=dtype)

b = a.clone().normal_()

res_add = torch.add(a, b, alpha=-1)

res_csub = a.clone()

res_csub.sub_(b)

self.assertEqual(res_add, res_csub)

# with a scalar

a = torch.randn(100, 100, dtype=dtype)

scalar = 123.5

res_add = torch.add(a, -scalar)

res_csub = a.clone()

res_csub.sub_(scalar)

self.assertEqual(res_add, res_csub)

@parametrize("dtype", list(get_all_math_dtypes("cpu")))

def test_rdiv(self, dtype):

if dtype is torch.float16:

return

elif dtype.is_complex:

x = torch.rand(100, dtype=dtype).add(1).mul(4)

else:

x = torch.rand(100).add(1).mul(4).to(dtype)

y = 30 / x

z = torch.tensor([30 / v.item() for v in x])

self.assertEqual(y, z, exact_dtype=False)

@parametrize("dtype", [torch.float])

def test_cdiv(self, dtype):

self._test_cop(torch.div, operator.truediv, dtype)

@parametrize("dtype", [torch.float])

def test_cremainder(self, dtype):

self._test_cop(torch.remainder, operator.mod, dtype)

@parametrize("dtype", [torch.float])

def test_cmul(self, dtype):

self._test_cop(torch.mul, operator.mul, dtype)

@parametrize("dtype", [torch.float])

def test_cpow(self, dtype):

self._test_cop(torch.pow, lambda x, y: nan if x < 0 else math.pow(x, y), dtype)

@parametrize(

"dtype", [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]

)

def test_floor_divide_zero(self, dtype):

a = torch.tensor([0, 1], dtype=dtype)

b = torch.tensor([0, 1], dtype=dtype)

with self.assertRaisesRegex(RuntimeError, "ZeroDivisionError"):

with self.assertWarnsOnceRegex(UserWarning, "floor_divide"):

a // b

@parametrize("dtype", [torch.half, torch.bfloat16])

def test_lerp_lowp_cpu(self, dtype):

xvals = (0.0, -30000.0)

yvals = (0.1, -20000.0)

for shape in [(4,), (20,), (3, 10, 10)]:

xs = [torch.full(shape, xval, dtype=dtype) for xval in xvals]

ys = [torch.full(shape, yval, dtype=dtype) for yval in yvals]

weights = [70000, torch.full(shape, 8, dtype=dtype)]

for x, y, w in zip(xs, ys, weights):

xref = x.float()

yref = y.float()

wref = w.float() if isinstance(w, torch.Tensor) else w

actual = torch.lerp(x, y, w)

expected = torch.lerp(xref, yref, wref).to(dtype)

self.assertEqual(actual, expected, atol=0.0, rtol=0.0)

class TestBinaryUfuncsDevice(TestCase):

# Generic tests for elementwise binary (AKA binary universal (u) functions (funcs))

# TODO: below contiguous tensor results are compared with a variety of noncontiguous results.

# It would be interesting to have the lhs and rhs have different discontinuities.

# Helper for comparing torch tensors and NumPy arrays

# TODO: should this or assertEqual also validate that strides are equal?

def assertEqualHelper(

self, actual, expected, msg, *, dtype, exact_dtype=True, **kwargs

):

if not isinstance(actual, torch.Tensor):

raise AssertionError(

f"expected actual to be torch.Tensor, got {type(actual)}"

)

# Some NumPy functions return scalars, not arrays

if isinstance(expected, Number):

self.assertEqual(actual.item(), expected, msg=msg, **kwargs)

elif isinstance(expected, np.ndarray):

# Handles exact dtype comparisons between arrays and tensors

if exact_dtype:

# Allows array dtype to be float32 when comparing with bfloat16 tensors

# since NumPy doesn't support the bfloat16 dtype

# Also ops like scipy.special.erf, scipy.special.erfc, etc, promote float16

# to float32

if expected.dtype == np.float32:

if actual.dtype not in (

torch.float16,

torch.bfloat16,

torch.float32,

):

raise AssertionError(

f"actual.dtype {actual.dtype} not in expected dtypes"

)

else:

if expected.dtype != torch_to_numpy_dtype_dict[actual.dtype]:

raise AssertionError(

f"dtype mismatch: {expected.dtype} != {torch_to_numpy_dtype_dict[actual.dtype]}"

)

self.assertEqual(

actual,

torch.from_numpy(expected).to(actual.dtype),

msg,

exact_device=False,

**kwargs,

)

else:

self.assertEqual(actual, expected, msg, exact_device=False, **kwargs)

# Tests that the function and its (array-accepting) reference produce the same

# values on given tensors

def _test_reference_numerics(self, dtype, op, gen, equal_nan=True):

def _helper_reference_numerics(

expected, actual, msg, exact_dtype, equal_nan=True

):

if not torch.can_cast(

numpy_to_torch_dtype_dict[expected.dtype.type], dtype

):

exact_dtype = False

if dtype is torch.bfloat16 and expected.dtype == np.float32:

# Ref: https://github.com/pytorch/pytorch/blob/master/torch/testing/_internal/common_utils.py#L1149

self.assertEqualHelper(

actual,

expected,

msg,

dtype=dtype,

exact_dtype=exact_dtype,

rtol=16e-3,

atol=1e-5,

)

else:

self.assertEqualHelper(

actual,

expected,

msg,

dtype=dtype,

equal_nan=equal_nan,

exact_dtype=exact_dtype,

)

for sample in gen:

# Each sample input acquired from the generator is just one lhs tensor

# and one rhs tensor

l = sample.input

r = sample.args[0]

numpy_sample = sample.numpy()

l_numpy = numpy_sample.input

r_numpy = numpy_sample.args[0]

actual = op(l, r)

expected = op.ref(l_numpy, r_numpy)

# Dtype promo rules have changed since NumPy 2.

# Specialize the backward-incompatible cases.

if (

np.__version__ > "2"

and op.name in ("sub", "_refs.sub")

and isinstance(l_numpy, np.ndarray)

):

expected = expected.astype(l_numpy.dtype)

# Crafts a custom error message for smaller, printable tensors

def _numel(x):

if isinstance(x, torch.Tensor):

return x.numel()

# Assumes x is a scalar

return 1

if _numel(l) <= 100 and _numel(r) <= 100:

msg = (

"Failed to produce expected results! Input lhs tensor was"

f" {l}, rhs tensor was {r}, torch result is {actual}, and reference result is"

f" {expected}."

)

else:

msg = None

exact_dtype = True

if isinstance(actual, torch.Tensor):

_helper_reference_numerics(

expected, actual, msg, exact_dtype, equal_nan

)

else:

for x, y in zip(expected, actual):

# testing multi-outputs results

_helper_reference_numerics(x, y, msg, exact_dtype, equal_nan)

# The following tests only apply to elementwise binary operators with references

binary_ufuncs_with_references = list(

filter(lambda op: op.ref is not None, binary_ufuncs)

)

@ops(binary_ufuncs_with_references)

def test_reference_numerics(self, device, dtype, op):

gen = generate_elementwise_binary_tensors(op, device=device, dtype=dtype)

self._test_reference_numerics(dtype, op, gen, equal_nan=True)

@ops(binary_ufuncs_with_references)

def test_reference_numerics_small_values(self, device, dtype, op):

if dtype is torch.bool:

self.skipTest("Doesn't support bool!")

gen = generate_elementwise_binary_small_value_tensors(

op, device=device, dtype=dtype

)

self._test_reference_numerics(dtype, op, gen, equal_nan=True)

@ops(

binary_ufuncs_with_references,

allowed_dtypes=(

torch.int16,

torch.int32,

torch.int64,

torch.float16,

torch.bfloat16,

torch.float32,

torch.float64,

torch.complex64,

torch.complex128,

),

)

def test_reference_numerics_large_values(self, device, dtype, op):

gen = generate_elementwise_binary_large_value_tensors(

op, device=device, dtype=dtype

)

self._test_reference_numerics(dtype, op, gen, equal_nan=True)

@ops(

binary_ufuncs_with_references,

allowed_dtypes=(

torch.float16,

torch.bfloat16,

torch.float32,

torch.float64,

torch.complex64,

torch.complex128,

),

)

def test_reference_numerics_extremal_values(self, device, dtype, op):

gen = generate_elementwise_binary_extremal_value_tensors(

op, device=device, dtype=dtype

)

self._test_reference_numerics(dtype, op, gen, equal_nan=True)

# tests broadcasting and noncontiguous broadcasting behavior

@ops(

binary_ufuncs_with_references,

allowed_dtypes=(

torch.long,

torch.float32,

),

)

def test_broadcasting(self, device, dtype, op):

gen = generate_elementwise_binary_broadcasting_tensors(

op, device=device, dtype=dtype

)

self._test_reference_numerics(dtype, op, gen, equal_nan=True)

@ops(

binary_ufuncs_with_references,

allowed_dtypes=(torch.long, torch.float32, torch.complex64),

)

def test_scalar_support(self, device, dtype, op):

gen = generate_elementwise_binary_with_scalar_samples(

op, device=device, dtype=dtype

)

self._test_reference_numerics(dtype, op, gen, equal_nan=True)

gen = generate_elementwise_binary_with_scalar_and_type_promotion_samples(

op, device=device, dtype=dtype

)

self._test_reference_numerics(dtype, op, gen, equal_nan=True)

@ops(binary_ufuncs)

def test_contig_vs_every_other(self, device, dtype, op):

lhs = make_tensor(

(1026,), device=device, dtype=dtype, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

(1026,), device=device, dtype=dtype, **op.rhs_make_tensor_kwargs

)

lhs_non_contig = lhs[::2]

rhs_non_contig = rhs[::2]

self.assertTrue(lhs.is_contiguous())

self.assertTrue(rhs.is_contiguous())

self.assertFalse(lhs_non_contig.is_contiguous())

self.assertFalse(rhs_non_contig.is_contiguous())

expected = op(lhs, rhs)[::2]

actual = op(lhs_non_contig, rhs_non_contig)

self.assertEqual(expected, actual)

@ops(binary_ufuncs)

def test_contig_vs_transposed(self, device, dtype, op):

lhs = make_tensor(

(789, 357), device=device, dtype=dtype, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

(789, 357), device=device, dtype=dtype, **op.rhs_make_tensor_kwargs

)

lhs_non_contig = lhs.T

rhs_non_contig = rhs.T

self.assertTrue(lhs.is_contiguous())

self.assertTrue(rhs.is_contiguous())

self.assertFalse(lhs_non_contig.is_contiguous())

self.assertFalse(rhs_non_contig.is_contiguous())

expected = op(lhs, rhs).T

actual = op(lhs_non_contig, rhs_non_contig)

self.assertEqual(expected, actual)

@ops(binary_ufuncs)

def test_non_contig(self, device, dtype, op):

shapes = ((5, 7), (1024,))

for shape in shapes:

lhs = make_tensor(

shape, dtype=dtype, device=device, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

shape, dtype=dtype, device=device, **op.rhs_make_tensor_kwargs

)

lhs_non_contig = torch.empty(shape + (2,), device=device, dtype=dtype)[

..., 0

]

lhs_non_contig.copy_(lhs)

rhs_non_contig = torch.empty(shape + (2,), device=device, dtype=dtype)[

..., 0

]

rhs_non_contig.copy_(rhs)

self.assertTrue(lhs.is_contiguous())

self.assertTrue(rhs.is_contiguous())

self.assertFalse(lhs_non_contig.is_contiguous())

self.assertFalse(rhs_non_contig.is_contiguous())

expected = op(lhs, rhs)

actual = op(lhs_non_contig, rhs_non_contig)

self.assertEqual(expected, actual)

@ops(binary_ufuncs)

def test_non_contig_index(self, device, dtype, op):

shape = (2, 2, 1, 2)

lhs = make_tensor(

shape, dtype=dtype, device=device, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

shape, dtype=dtype, device=device, **op.rhs_make_tensor_kwargs

)

lhs_non_contig = lhs[:, 1, ...]

lhs = lhs_non_contig.contiguous()

rhs_non_contig = rhs[:, 1, ...]

rhs = rhs_non_contig.contiguous()

self.assertTrue(lhs.is_contiguous())

self.assertTrue(rhs.is_contiguous())

self.assertFalse(lhs_non_contig.is_contiguous())

self.assertFalse(rhs_non_contig.is_contiguous())

expected = op(lhs, rhs)

actual = op(lhs_non_contig, rhs_non_contig)

self.assertEqual(expected, actual)

@ops(binary_ufuncs)

def test_non_contig_expand(self, device, dtype, op):

shapes = [(1, 3), (1, 7), (5, 7)]

for shape in shapes:

lhs = make_tensor(

shape, dtype=dtype, device=device, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

shape, dtype=dtype, device=device, **op.rhs_make_tensor_kwargs

)

lhs_non_contig = lhs.clone().expand(3, -1, -1)

rhs_non_contig = rhs.clone().expand(3, -1, -1)

self.assertTrue(lhs.is_contiguous())

self.assertTrue(rhs.is_contiguous())

self.assertFalse(lhs_non_contig.is_contiguous())

self.assertFalse(rhs_non_contig.is_contiguous())

expected = op(lhs, rhs)

actual = op(lhs_non_contig, rhs_non_contig)

for i in range(3):

self.assertEqual(expected, actual[i])

@ops(binary_ufuncs)

def test_contig_size1(self, device, dtype, op):

shape = (5, 100)

lhs = make_tensor(

shape, dtype=dtype, device=device, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

shape, dtype=dtype, device=device, **op.rhs_make_tensor_kwargs

)

lhs = lhs[:1, :50]

lhs_alt = torch.empty(lhs.size(), device=device, dtype=dtype)

lhs_alt.copy_(lhs)

rhs = rhs[:1, :50]

rhs_alt = torch.empty(rhs.size(), device=device, dtype=dtype)

rhs_alt.copy_(rhs)

self.assertTrue(lhs.is_contiguous())

self.assertTrue(rhs.is_contiguous())

self.assertTrue(lhs_alt.is_contiguous())

self.assertTrue(rhs_alt.is_contiguous())

expected = op(lhs, rhs)

actual = op(lhs_alt, rhs_alt)

self.assertEqual(expected, actual)

@ops(binary_ufuncs)

def test_contig_size1_large_dim(self, device, dtype, op):

shape = (5, 2, 3, 1, 4, 5, 3, 2, 1, 2, 3, 4)

lhs = make_tensor(

shape, dtype=dtype, device=device, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

shape, dtype=dtype, device=device, **op.rhs_make_tensor_kwargs

)

lhs = lhs[:1, :, :, :, :, :, :, :, :, :, :, :]

lhs_alt = torch.empty(lhs.size(), device=device, dtype=dtype)

lhs_alt.copy_(lhs)

rhs = rhs[:1, :, :, :, :, :, :, :, :, :, :, :]

rhs_alt = torch.empty(rhs.size(), device=device, dtype=dtype)

rhs_alt.copy_(rhs)

self.assertTrue(lhs.is_contiguous())

self.assertTrue(rhs.is_contiguous())

self.assertTrue(lhs_alt.is_contiguous())

self.assertTrue(rhs_alt.is_contiguous())

expected = op(lhs, rhs)

actual = op(lhs_alt, rhs_alt)

self.assertEqual(expected, actual)

@ops(binary_ufuncs)

def test_batch_vs_slicing(self, device, dtype, op):

shape = (32, 512)

lhs = make_tensor(

shape, dtype=dtype, device=device, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

shape, dtype=dtype, device=device, **op.rhs_make_tensor_kwargs

)

expected = op(lhs, rhs)

actual = []

for idx in range(32):

actual.append(op(lhs[idx], rhs[idx]))

actual = torch.stack(actual)

self.assertEqual(expected, actual)

# Tests that elementwise binary operators participate in type promotion properly

# NOTE: because the cross-product of all possible type promotion tests is huge, this

# just spot checks some handwritten cases.

# NOTE: It may be possible to refactor this test into something simpler

@ops(binary_ufuncs, dtypes=OpDTypes.none)

def test_type_promotion(self, device, op):

supported_dtypes = op.supported_dtypes(torch.device(device).type)

make_lhs = partial(

make_tensor, (5,), device=device, **op.lhs_make_tensor_kwargs

)

make_rhs = partial(

make_tensor, (5,), device=device, **op.rhs_make_tensor_kwargs

)

make_rhs_scalar_tensor = partial(

make_tensor, (), device="cpu", **op.rhs_make_tensor_kwargs

)

def _supported(dtypes):

return all(x in supported_dtypes for x in dtypes)

# int x int type promotion

if _supported((torch.int16, torch.int32, torch.int64)):

lhs_i16 = make_lhs(dtype=torch.int16)

lhs_i32 = make_lhs(dtype=torch.int32)

lhs_i64 = make_lhs(dtype=torch.int64)

rhs_i16 = make_rhs(dtype=torch.int16)

rhs_i32 = make_rhs(dtype=torch.int32)

rhs_i64 = make_rhs(dtype=torch.int64)

if op.promotes_int_to_float:

default_dtype = torch.get_default_dtype()

self.assertEqual(op(lhs_i16, rhs_i32).dtype, default_dtype)

self.assertEqual(

op(lhs_i16, rhs_i32),

op(lhs_i16.to(default_dtype), rhs_i32.to(default_dtype)),

)

self.assertEqual(op(lhs_i32, rhs_i64).dtype, default_dtype)

self.assertEqual(

op(lhs_i32, rhs_i64),

op(lhs_i32.to(default_dtype), rhs_i64.to(default_dtype)),

)

elif op.always_returns_bool:

self.assertEqual(op(lhs_i16, rhs_i32).dtype, torch.bool)

self.assertEqual(op(lhs_i32, rhs_i64).dtype, torch.bool)

else: # standard type promotion

self.assertEqual(op(lhs_i16, rhs_i32).dtype, torch.int32)

self.assertEqual(

op(lhs_i16, rhs_i32), op(lhs_i16.to(torch.int32), rhs_i32)

)

self.assertEqual(op(lhs_i32, rhs_i64).dtype, torch.int64)

self.assertEqual(

op(lhs_i32, rhs_i64), op(lhs_i32.to(torch.int64), rhs_i64)

)

if op.supports_out:

if not op.promotes_int_to_float:

# Integers can be safely cast to other integer types

out = torch.empty_like(lhs_i64)

self.assertEqual(op(lhs_i16, rhs_i32, out=out).dtype, torch.int64)

self.assertEqual(op(lhs_i16, rhs_i32), out, exact_dtype=False)

out = torch.empty_like(lhs_i16)

self.assertEqual(op(lhs_i32, rhs_i64, out=out).dtype, torch.int16)

else:

# Float outs cannot be safely cast to integer types

with self.assertRaisesRegex(RuntimeError, "can't be cast"):

op(lhs_i16, rhs_i32, out=torch.empty_like(lhs_i64))

if not op.always_returns_bool:

# Neither integer nor float outs can be cast to bool

with self.assertRaisesRegex(RuntimeError, "can't be cast"):

op(

lhs_i16,

rhs_i32,

out=torch.empty_like(lhs_i64, dtype=torch.bool),

)

# All these output types can be cast to any float or complex type

out = torch.empty_like(lhs_i64, dtype=torch.float16)

self.assertEqual(op(lhs_i16, rhs_i32, out=out).dtype, torch.float16)

out = torch.empty_like(lhs_i64, dtype=torch.bfloat16)

self.assertEqual(op(lhs_i16, rhs_i32, out=out).dtype, torch.bfloat16)

out = torch.empty_like(lhs_i64, dtype=torch.float32)

self.assertEqual(op(lhs_i16, rhs_i32, out=out).dtype, torch.float32)

self.assertEqual(op(lhs_i16, rhs_i32), out, exact_dtype=False)

out = torch.empty_like(lhs_i64, dtype=torch.complex64)

self.assertEqual(op(lhs_i16, rhs_i32, out=out).dtype, torch.complex64)

self.assertEqual(op(lhs_i16, rhs_i32), out, exact_dtype=False)

# float x float type promotion

if _supported((torch.float32, torch.float64)):

lhs_f32 = make_lhs(dtype=torch.float32)

lhs_f64 = make_lhs(dtype=torch.float64)

rhs_f32 = make_rhs(dtype=torch.float32)

rhs_f64 = make_rhs(dtype=torch.float64)

if op.always_returns_bool:

self.assertEqual(op(lhs_f32, rhs_f64).dtype, torch.bool)

else: # normal float type promotion

self.assertEqual(op(lhs_f32, rhs_f64).dtype, torch.float64)

self.assertEqual(

op(lhs_f32, rhs_f64), op(lhs_f32.to(torch.float64), rhs_f64)

)

if op.supports_out:

# All these output types can be cast to any float or complex type

out = torch.empty_like(lhs_f64, dtype=torch.float16)

self.assertEqual(op(lhs_f32, rhs_f64, out=out).dtype, torch.float16)

out = torch.empty_like(lhs_f64, dtype=torch.bfloat16)

self.assertEqual(op(lhs_f32, rhs_f64, out=out).dtype, torch.bfloat16)

self.assertEqual(op(lhs_f32, rhs_f64), out, exact_dtype=False)

out = torch.empty_like(lhs_f64, dtype=torch.float32)

self.assertEqual(op(lhs_f32, rhs_f64, out=out).dtype, torch.float32)

self.assertEqual(op(lhs_f32, rhs_f64), out, exact_dtype=False)

out = torch.empty_like(lhs_f64, dtype=torch.complex64)

self.assertEqual(op(lhs_f32, rhs_f64, out=out).dtype, torch.complex64)

self.assertEqual(op(lhs_f32, rhs_f64), out, exact_dtype=False)

if not op.always_returns_bool:

# float outs can't be cast to an integer dtype

with self.assertRaisesRegex(RuntimeError, "can't be cast"):

op(

lhs_f32,

rhs_f64,

out=torch.empty_like(lhs_f64, dtype=torch.int64),

)

else:

# bool outs can be cast to an integer dtype

out = torch.empty_like(lhs_f64, dtype=torch.int64)

self.assertEqual(op(lhs_f32, rhs_f64, out=out).dtype, torch.int64)

self.assertEqual(op(lhs_f32, rhs_f64), out, exact_dtype=False)

# complex x complex type promotion

if _supported((torch.complex64, torch.complex128)):

lhs_c64 = make_lhs(dtype=torch.complex64)

lhs_c128 = make_lhs(dtype=torch.complex128)

rhs_c64 = make_rhs(dtype=torch.complex64)

rhs_c128 = make_rhs(dtype=torch.complex128)

if op.always_returns_bool:

self.assertEqual(op(lhs_c64, lhs_c128).dtype, torch.bool)

else: # normal complex type promotion

self.assertEqual(op(lhs_c64, rhs_c128).dtype, torch.complex128)

self.assertEqual(

op(lhs_c64, rhs_c128), op(lhs_c64.to(torch.complex128), rhs_c128)

)

if op.supports_out:

# All these output types can be cast to any or complex type

out = torch.empty_like(lhs_c64, dtype=torch.complex64)

self.assertEqual(op(lhs_c64, rhs_c128, out=out).dtype, torch.complex64)

result = op(lhs_c64, rhs_c128)

self.assertEqual(result, out.to(result.dtype))

if not op.always_returns_bool:

# complex outs can't be cast to float types

with self.assertRaisesRegex(RuntimeError, "can't be cast"):

op(

lhs_c64,

rhs_c128,

out=torch.empty_like(lhs_c64, dtype=torch.float64),

)

# complex outs can't be cast to an integer dtype

with self.assertRaisesRegex(RuntimeError, "can't be cast"):

op(

lhs_c64,

rhs_c128,

out=torch.empty_like(lhs_c64, dtype=torch.int64),

)

else:

# bool outs can be cast to a float type

out = torch.empty_like(lhs_c64, dtype=torch.float64)

self.assertEqual(

op(lhs_c64, rhs_c128, out=out).dtype, torch.float64

)

self.assertEqual(op(lhs_c64, rhs_c128), out, exact_dtype=False)

# bool outs can be cast to an integer dtype

out = torch.empty_like(lhs_f64, dtype=torch.int64)

self.assertEqual(op(lhs_f32, rhs_f64, out=out).dtype, torch.int64)

self.assertEqual(op(lhs_f32, rhs_f64), out, exact_dtype=False)

# int x float type promotion

# Note: float type is the result dtype

if _supported((torch.long, torch.float32)):

lhs_i64 = make_lhs(dtype=torch.int64)

rhs_f32 = make_rhs(dtype=torch.float32)

result = op(lhs_i64, rhs_f32)

expected_dtype = torch.float32 if not op.always_returns_bool else torch.bool

self.assertEqual(result.dtype, expected_dtype)

# float x complex type promotion

# Note: complex type with highest "value type" is the result dtype

if _supported((torch.float64, torch.complex64)):

lhs_f64 = make_lhs(dtype=torch.float64)

rhs_c64 = make_rhs(dtype=torch.complex64)

result = op(lhs_f64, rhs_c64)

expected_dtype = (

torch.complex128 if not op.always_returns_bool else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# int x float scalar type promotion

# Note: default float dtype is the result dtype

if _supported((torch.int64, torch.float32)) and op.supports_rhs_python_scalar:

lhs_i64 = make_lhs(dtype=torch.int64)

rhs_f_scalar = 1.0

result = op(lhs_i64, rhs_f_scalar)

expected_dtype = (

torch.get_default_dtype() if not op.always_returns_bool else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# repeats with a scalar float tensor, which should set the dtype

rhs_f32_scalar_tensor = make_rhs_scalar_tensor(dtype=torch.float32)

result = op(lhs_i64, rhs_f32_scalar_tensor)

expected_dtype = torch.float32 if not op.always_returns_bool else torch.bool

self.assertEqual(result.dtype, expected_dtype)

# Additional test with double

if _supported((torch.float64,)):

rhs_f64_scalar_tensor = make_rhs_scalar_tensor(dtype=torch.float64)

result = op(lhs_i64, rhs_f64_scalar_tensor)

expected_dtype = (

torch.float64 if not op.always_returns_bool else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# float x complex scalar type promotion

# Note: result dtype is complex with highest "value type" among all tensors

if (

_supported((torch.float32, torch.complex64))

and op.supports_rhs_python_scalar

):

lhs_f32 = make_lhs(dtype=torch.float32)

rhs_c_scalar = complex(1, 1)

result = op(lhs_f32, rhs_c_scalar)

expected_dtype = (

torch.complex64 if not op.always_returns_bool else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# repeats with a scalar complex tensor

rhs_c64_scalar_tensor = make_rhs_scalar_tensor(dtype=torch.complex64)

result = op(lhs_f32, rhs_c64_scalar_tensor)

expected_dtype = (

torch.complex64 if not op.always_returns_bool else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# Additional test with complexdouble

if _supported((torch.complex128,)):

rhs_c128_scalar_tensor = make_rhs_scalar_tensor(dtype=torch.complex128)

result = op(lhs_f32, rhs_c128_scalar_tensor)

# Value type of 1D+ Tensor (lhs_f32) takes priority over scalar tensor (rhs_c128).

expected_dtype = (

torch.complex64 if not op.always_returns_bool else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# float x float scalar tensor

# Note: result dtype is the type of the float tensor

if _supported((torch.float32, torch.float64)) and op.supports_rhs_python_scalar:

lhs_f32 = make_lhs(dtype=torch.float32)

rhs_f64_scalar_tensor = make_rhs_scalar_tensor(dtype=torch.float64)

result = op(lhs_f32, rhs_f64_scalar_tensor)

expected_dtype = torch.float32 if not op.always_returns_bool else torch.bool

self.assertEqual(result.dtype, expected_dtype)

# complex x complex scalar tensor

# Note: result dtype is the type of the complex tensor

if (

_supported((torch.complex64, torch.complex128))

and op.supports_rhs_python_scalar

):

lhs_c64 = make_lhs(dtype=torch.complex64)

rhs_c128_scalar_tensor = make_rhs_scalar_tensor(dtype=torch.complex128)

result = op(lhs_c64, rhs_c128_scalar_tensor)

expected_dtype = (

torch.complex64 if not op.always_returns_bool else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# scalar x scalar

# Note: result dtype is default float type

if op.supports_two_python_scalars and _supported((torch.long, torch.float32)):

rhs_f_scalar = 2.0

for lhs in (1, 1.0):

result = op(lhs, rhs_f_scalar)

expected_dtype = (

torch.get_default_dtype()

if not op.always_returns_bool

else torch.bool

)

self.assertEqual(result.dtype, expected_dtype)

# TODO: move to error input test

@ops(binary_ufuncs, allowed_dtypes=(torch.float32,))

def test_not_broadcastable(self, device, dtype, op):

for shape_lhs, shape_rhs in (

((2,), (3,)),

((3, 1), (2, 1)),

((1, 3, 2), (3,)),

((3, 1, 2), (2, 1, 2)),

):

lhs = make_tensor(

shape_lhs, device=device, dtype=dtype, **op.lhs_make_tensor_kwargs

)

rhs = make_tensor(

shape_rhs, device=device, dtype=dtype, **op.rhs_make_tensor_kwargs

)

try:

broadcasted_shape = op(lhs, rhs).shape

except RuntimeError:

continue

msg = (

f"On {device}, torch.{op.name} broadcasts inputs shapes {shape_lhs} and {shape_rhs} into "

f"{broadcasted_shape}, although they are not broadcastable."

)

raise AssertionError(msg)

def test_add_broadcast_empty(self, device):

# empty + empty

self.assertRaises(

RuntimeError,

lambda: torch.randn(5, 0, device=device) + torch.randn(0, 5, device=device),

)

self.assertEqual(

torch.randn(5, 0, device=device),

torch.randn(0, device=device) + torch.randn(5, 0, device=device),

)

self.assertEqual(

torch.randn(5, 0, 0, device=device),

torch.randn(0, device=device) + torch.randn(5, 0, 1, device=device),

)