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)