def test_bfloat16(self, device):
# with scalar
bf = torch.tensor(5.5, dtype=torch.bfloat16, device=device)
for scalar in (2.2, 5, 100000): # bf + 100000 is inf
self.assertEqual((bf + scalar).dtype, torch.bfloat16)
self.assertEqual(scalar + bf, bf + scalar)
for scalar in (complex(1, 1), complex(-2, 0), complex(0, -3)):
self.assertEqual((bf + scalar).dtype, torch.cfloat)
self.assertEqual(bf + scalar, scalar + bf)
# with tensor
for dtype in torch.testing.get_all_dtypes():
t = torch.tensor(1, dtype=dtype, device=device)
self.assertEqual(bf + t, t + bf)
if dtype in (torch.float16, torch.float32, torch.float64, torch.cfloat, torch.cdouble):
# Handles bfloat16 x float16 -> float32 promotion
expected_dtype = dtype if dtype != torch.half else torch.float32
elif dtype in (torch.bool, torch.uint8,
torch.int8, torch.int16, torch.int32, torch.int64, torch.bfloat16):
expected_dtype = torch.bfloat16
else:
raise AssertionError(f'Missing dtype {dtype} not tested.')
self.assertEqual(torch.promote_types(dtype, torch.bfloat16), expected_dtype)
self.assertEqual(torch.promote_types(torch.bfloat16, dtype), expected_dtype)
self.assertEqual((bf + t).dtype, expected_dtype)
def test_promote_types(self, device):
self.assertEqual(torch.promote_types(torch.float, torch.int),
torch.float)
self.assertEqual(torch.promote_types(torch.float, torch.double),
torch.double)
self.assertEqual(torch.promote_types(torch.int, torch.uint8),
torch.int)
def test_result_type(self, device, dtypes):
"Test result_type for tensor vs tensor and scalar vs scalar."
def _get_dtype(x):
"Get the dtype of x if x is a tensor. If x is a scalar, get its corresponding dtype if it were a tensor."
if torch.is_tensor(x):
return x.dtype
elif isinstance(x, bool):
return torch.bool
elif isinstance(x, int):
return torch.int64
elif isinstance(x, float):
return torch.float32
elif isinstance(x, complex):
return torch.complex64
else:
raise AssertionError(f"Unkonwn type {x}")
# tensor against tensor
a_tensor = torch.tensor((0, 1), device=device, dtype=dtypes[0])
a_single_tensor = torch.tensor(1, device=device, dtype=dtypes[0])
a_scalar = a_single_tensor.item()
b_tensor = torch.tensor((1, 0), device=device, dtype=dtypes[1])
b_single_tensor = torch.tensor(1, device=device, dtype=dtypes[1])
b_scalar = b_single_tensor.item()
combo = ((a_tensor, a_single_tensor, a_scalar),
(b_tensor, b_single_tensor, b_scalar))
for a, b in itertools.product(*combo):
dtype_a = _get_dtype(a)
dtype_b = _get_dtype(b)
try:
result = a + b
except RuntimeError:
with self.assertRaises(RuntimeError):
torch.promote_types(dtype_a, dtype_b)
with self.assertRaises(RuntimeError):
torch.result_type(a, b)
else:
dtype_res = _get_dtype(result)
if a is a_scalar and b is b_scalar and dtype_a == torch.bool and dtype_b == torch.bool:
# special case: in Python, True + True is an integer
self.assertEqual(dtype_res, torch.int64,
f"a == {a}, b == {b}")
else:
self.assertEqual(dtype_res, torch.result_type(a, b),
f"a == {a}, b == {b}")
if a is a_scalar and b is b_scalar: # Python internal type determination is good enough in this case
continue
if any(a is a0 and b is b0 for a0, b0 in zip(
*combo)): # a and b belong to the same class
self.assertEqual(dtype_res,
torch.promote_types(dtype_a, dtype_b),
f"a == {a}, b == {b}")
def test_many_promotions(self, device):
# Can also include half on CPU in cases where it will be promoted to a
# supported dtype
dtypes1 = torch.testing.get_all_math_dtypes('cuda')
dtypes2 = torch.testing.get_all_math_dtypes(device)
ops = [torch.add, torch.sub, torch.mul, torch.div, torch.rsub]
for dt1, dt2 in itertools.product(dtypes1, dtypes2):
for op, non_contiguous in itertools.product(ops, [True, False]):
common_dtype = torch.promote_types(dt1, dt2)
if common_dtype == torch.half and self.device_type == 'cpu':
continue
if op == torch.sub and common_dtype != torch.bool:
# Subtraction, the `-` operator, with a bool tensor is not supported.
continue
first = self._get_test_tensor(device, dt1)
second = self._get_test_tensor(device, dt2, op == torch.div)
# test ops with non-contiguous tensors
if non_contiguous:
first = first.transpose(0, 2)
second = second.transpose(2, 1)
self.assertNotEqual(first.stride(), second.stride(),
msg="some non-contiguous issues could be missed if tensors have same strides")
self.assertEqual(not first.is_contiguous(), non_contiguous)
self.assertEqual(not second.is_contiguous(), non_contiguous)
result = op(first, second)
expected = op(first.to(common_dtype), second.to(common_dtype))
self.assertEqual(result.dtype, expected.dtype, msg='{} with {}, {}'.format(op.__name__, dt1, dt2))
self.assertEqual(result, expected, msg='{} with {}, {}'.format(op.__name__, dt1, dt2))
def _equalize_attributes(actual: Tensor,
expected: Tensor) -> Tuple[Tensor, Tensor]:
"""Equalizes some attributes of two tensors for value comparison.
If :attr:`actual` and :attr:`expected`
- are not onn the same memory :attr:`~torch.Tensor.device`, they are moved CPU memory, and
- do not have the same :attr:`~torch.Tensor.dtype`, they are copied to the :class:`~torch.dtype` returned by
:func:`torch.promote_types`.
Args:
actual (Tensor): Actual tensor.
expected (Tensor): Expected tensor.
Returns:
Tuple(Tensor, Tensor): Equalized tensors.
"""
if actual.device != expected.device:
actual = actual.cpu()
expected = expected.cpu()
if actual.dtype != expected.dtype:
dtype = torch.promote_types(actual.dtype, expected.dtype)
actual = actual.to(dtype)
expected = expected.to(dtype)
return actual, expected
def __init__(self, func, y0, rtol, atol, first_step=None, safety=0.9, ifactor=10.0, dfactor=0.2,
max_num_steps=2 ** 31 - 1, grid_points=None, eps=0., dtype=torch.float64, **kwargs):
super(RKAdaptiveStepsizeODESolver, self).__init__(dtype=dtype, y0=y0, **kwargs)
# We use mixed precision. y has its original dtype (probably float32), whilst all 'time'-like objects use
# `dtype` (defaulting to float64).
dtype = torch.promote_types(dtype, y0.dtype)
device = y0.device
self.func = lambda t, y: func(t.type_as(y), y)
self.rtol = torch.as_tensor(rtol, dtype=dtype, device=device)
self.atol = torch.as_tensor(atol, dtype=dtype, device=device)
self.first_step = None if first_step is None else torch.as_tensor(first_step, dtype=dtype, device=device)
self.safety = torch.as_tensor(safety, dtype=dtype, device=device)
self.ifactor = torch.as_tensor(ifactor, dtype=dtype, device=device)
self.dfactor = torch.as_tensor(dfactor, dtype=dtype, device=device)
self.max_num_steps = torch.as_tensor(max_num_steps, dtype=torch.int32, device=device)
grid_points = torch.tensor([], dtype=dtype, device=device) if grid_points is None else grid_points.to(dtype)
self.grid_points = grid_points
self.eps = torch.as_tensor(eps, dtype=dtype, device=device)
self.dtype = dtype
# Copy from class to instance to set device
self.tableau = _ButcherTableau(alpha=self.tableau.alpha.to(device=device, dtype=y0.dtype),
beta=[b.to(device=device, dtype=y0.dtype) for b in self.tableau.beta],
c_sol=self.tableau.c_sol.to(device=device, dtype=y0.dtype),
c_error=self.tableau.c_error.to(device=device, dtype=y0.dtype))
self.mid = self.mid.to(device=device, dtype=y0.dtype)
def _equalize_attributes(a: torch.Tensor,
b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""Equalizes some attributes of two tensors for value comparison.
If :attr:`a` and :attr:`b`
- do not live in the same memory :attr:`~torch.Tensor.device`, they are moved CPU memory, and
- do not have the same :attr:`~torch.Tensor.dtype`, they are copied to the :class:`~torch.dtype` returned by
:func:`torch.promote_types`.
Args:
a (torch.Tensor): First tensor.
b (torch.Tensor): Second tensor.
Returns:
Tuple(torch.Tensor, torch.Tensor): Equalized tensors.
"""
if a.device != b.device:
a = a.cpu()
b = b.cpu()
if a.dtype != b.dtype:
dtype = torch.promote_types(a.dtype, b.dtype)
a = a.to(dtype)
b = b.to(dtype)
return a, b
def _equalize_attributes(actual: Tensor, expected: Tensor) -> Tuple[Tensor, Tensor]:
"""Equalizes some attributes of two tensors for value comparison.
If :attr:`actual` and :attr:`expected`
- are not on the same :attr:`~torch.Tensor.device`, they are moved CPU memory, and
- do not have the same ``dtype``, they are promoted to a common ``dtype`` (according to
:func:`torch.promote_types`)
Args:
actual (Tensor): Actual tensor.
expected (Tensor): Expected tensor.
Returns:
Tuple(Tensor, Tensor): Equalized tensors.
"""
if actual.device != expected.device:
actual = actual.cpu()
expected = expected.cpu()
if actual.dtype != expected.dtype:
dtype = torch.promote_types(actual.dtype, expected.dtype)
actual = actual.to(dtype)
expected = expected.to(dtype)
if actual.is_sparse and actual.is_coalesced() != expected.is_coalesced():
actual = actual.coalesce()
expected = expected.coalesce()
return actual, expected
def _equalize_attributes(actual: Tensor,
expected: Tensor) -> Tuple[Tensor, Tensor]:
"""Equalizes some attributes of two tensors for value comparison.
If :attr:`actual` and :attr:`expected` are ...
- ... not on the same :attr:`~torch.Tensor.device`, they are moved CPU memory.
- ... not of the same ``dtype``, they are promoted to a common ``dtype`` (according to
:func:`torch.promote_types`).
- ... not of the same ``layout``, they are converted to strided tensors.
- ... both sparse COO tensors but only one is coalesced, the other one is coalesced.
Args:
actual (Tensor): Actual tensor.
expected (Tensor): Expected tensor.
Returns:
Tuple(Tensor, Tensor): Equalized tensors.
"""
if actual.device != expected.device:
actual = actual.cpu()
expected = expected.cpu()
if actual.dtype != expected.dtype:
dtype = torch.promote_types(actual.dtype, expected.dtype)
actual = actual.to(dtype)
expected = expected.to(dtype)
if actual.layout != expected.layout:
# These checks are needed, since Tensor.to_dense() fails on tensors that are already strided
if actual.layout != torch.strided:
actual = actual.to_dense()
if expected.layout != torch.strided:
expected = expected.to_dense()
elif actual.is_sparse and actual.is_coalesced() != expected.is_coalesced():
actual = actual.coalesce()
expected = expected.coalesce()
return actual, expected
def test_promote_self(self, device):
for dtype in torch.testing.get_all_dtypes():
self.assertEqual(torch.promote_types(dtype, dtype), dtype)
def _get_default_rtol_and_atol(actual: Tensor,
expected: Tensor) -> Tuple[float, float]:
dtype = actual.dtype if actual.dtype == expected.dtype else torch.promote_types(
actual.dtype, expected.dtype)
return _DTYPE_PRECISIONS.get(dtype, (0.0, 0.0))
def _get_default_rtol_and_atol(a: torch.Tensor,
b: torch.Tensor) -> Tuple[float, float]:
dtype = a.dtype if a.dtype == b.dtype else torch.promote_types(
a.dtype, b.dtype)
return _DTYPE_PRECISIONS.get(dtype, (0.0, 0.0))
