def test_min_max(self, device, dtype):
for N in N_values:
tensors1 = self._get_test_data(device, dtype, N)
tensors2 = self._get_test_data(device, dtype, N)
# Mimics cuda kernel dtype flow. With fp16/bf16 input, runs in fp32 and casts output back to fp16/bf16.
control_dtype = torch.float32 if (
self.device_type == 'cuda' and
(dtype is torch.float16 or dtype is torch.bfloat16)) else dtype
expected_max = [
torch.max(tensors1[i].to(dtype=control_dtype),
tensors2[i].to(dtype=control_dtype)).to(dtype=dtype)
for i in range(N)
]
expected_min = [
torch.min(tensors1[i].to(dtype=control_dtype),
tensors2[i].to(dtype=control_dtype)).to(dtype=dtype)
for i in range(N)
]
res_max = torch._foreach_maximum(tensors1, tensors2)
self.assertEqual(res_max, expected_max)
res_min = torch._foreach_minimum(tensors1, tensors2)
self.assertEqual(res_min, expected_min)
def test_max_min_inf_nan(self, device, dtype):
a = [
torch.tensor([inf], device=device, dtype=dtype),
torch.tensor([-inf], device=device, dtype=dtype),
torch.tensor([nan], device=device, dtype=dtype),
torch.tensor([nan], device=device, dtype=dtype)
]
b = [
torch.tensor([-inf], device=device, dtype=dtype),
torch.tensor([inf], device=device, dtype=dtype),
torch.tensor([inf], device=device, dtype=dtype),
torch.tensor([nan], device=device, dtype=dtype)
]
expected_max = [torch.max(a1, b1) for a1, b1 in zip(a, b)]
res_max = torch._foreach_maximum(a, b)
self.assertEqual(expected_max, res_max)
expected_min = [torch.min(a1, b1) for a1, b1 in zip(a, b)]
res_min = torch._foreach_minimum(a, b)
self.assertEqual(expected_min, res_min)
def test_max_min_float_inf_nan(self, device, dtype):
a = [
torch.tensor([float('inf')], device=device, dtype=dtype),
torch.tensor([-float('inf')], device=device, dtype=dtype),
torch.tensor([float('nan')], device=device, dtype=dtype),
torch.tensor([float('nan')], device=device, dtype=dtype)
]
b = [
torch.tensor([-float('inf')], device=device, dtype=dtype),
torch.tensor([float('inf')], device=device, dtype=dtype),
torch.tensor([float('inf')], device=device, dtype=dtype),
torch.tensor([float('nan')], device=device, dtype=dtype)
]
expected = [torch.max(a1, b1) for a1, b1 in zip(a, b)]
res = torch._foreach_maximum(a, b)
self.assertEqual(expected, res)
expected = [torch.min(a1, b1) for a1, b1 in zip(a, b)]
res = torch._foreach_minimum(a, b)
self.assertEqual(expected, res)
