def test_autograd_to_mkldnn(self):
# MKLDNN only supports float32
root = torch.randn(4, 5, dtype=torch.float32, requires_grad=True)
def func(root):
return root.to_mkldnn().to_dense()
# because MKLDNN only supports float32, we need to lessen the precision.
# these numbers are just empirical results that seem to work.
self.assertWarnsRegex(UserWarning,
'double precision floating point',
lambda: gradcheck(func, [root], atol=4e-2, rtol=1e-2))
self.assertWarnsRegex(UserWarning,
'double precision floating point',
lambda: gradgradcheck(func, [root], atol=4e-2, rtol=1e-2))
def _check_helper(self, device, dtype, op, variant, check, *, check_forward_ad=False, check_backward_ad=True,
check_batched_grad=None, check_batched_forward_grad=False):
assert check in ('gradcheck', 'bwgrad_bwgrad', 'fwgrad_bwgrad')
# NB: check_backward_ad does not affect gradgradcheck (always True)
if variant is None:
self.skipTest("Skipped! Variant not implemented.")
if not op.supports_dtype(dtype, torch.device(device).type):
self.skipTest(f"Skipped! {op.name} does not support dtype {str(dtype)}")
def is_inplace(variant):
if hasattr(variant, "__wrapped__"):
return variant.__wrapped__ is op.get_inplace()
return variant is op.get_inplace()
include_conjugated_inputs = op.test_conjugated_samples and dtype.is_complex
samples = op.sample_inputs(device, dtype, requires_grad=True, include_conjugated_inputs=include_conjugated_inputs,
small_inputs_only=is_slow_gradcheck_env())
for sample in samples:
if sample.broadcasts_input and is_inplace(variant):
continue
# Gradcheck expects tensors as its input, but autograd actually supports tensorlists
# and tensors passed as kwargs. The following creates a function that accepts just
# the tensors that require grad as varargs, and then recomposes them back into the
# original input.
# Creates gradcheck inputs by identifying tensors requiring grad
all_args = None
if is_iterable_of_tensors(sample.input):
all_args = chain(sample.input, sample.args, sample.kwargs.values())
else:
all_args = tuple(chain((sample.input,), sample.args, sample.kwargs.values()))
gradcheck_args = tuple(x for x in all_args if (isinstance(x, torch.Tensor) and x.requires_grad))
def _input_recomposition_helper(inputs, inp, input_idx):
if is_iterable_of_tensors(inp):
tensor_list = []
for x in inp:
if isinstance(x, torch.Tensor) and x.requires_grad:
tensor_list.append(inputs[input_idx])
input_idx = input_idx + 1
else:
tensor_list.append(x)
return tensor_list, input_idx
elif isinstance(inp, torch.Tensor) and inp.requires_grad:
return inputs[input_idx], input_idx + 1
else:
return inp, input_idx
def fn(*inputs):
# Puts inputs back into sample properly
positional_args = []
input_idx = 0
inp, input_idx = _input_recomposition_helper(inputs, sample.input, input_idx)
positional_args.append(inp)
for x in sample.args:
inp, input_idx = _input_recomposition_helper(inputs, x, input_idx)
positional_args.append(inp)
# Recreates kwargs
kwargs = {}
for k, v in sample.kwargs.items():
inp, input_idx = _input_recomposition_helper(inputs, v, input_idx)
kwargs[k] = inp
output = op.gradcheck_wrapper(variant, *positional_args, **kwargs)
if sample.output_process_fn_grad is not None:
return sample.output_process_fn_grad(output)
return output
if check == 'gradcheck':
if check_batched_grad is None:
check_batched_grad = op.check_batched_grad
self.assertTrue(gradcheck(fn, gradcheck_args,
check_batched_grad=check_batched_grad,
check_grad_dtypes=True,
nondet_tol=op.gradcheck_nondet_tol,
fast_mode=op.gradcheck_fast_mode,
check_forward_ad=check_forward_ad,
check_backward_ad=check_backward_ad,
check_undefined_grad=True,
check_batched_forward_grad=check_batched_forward_grad))
elif check in ('bwgrad_bwgrad', 'fwgrad_bwgrad'): # gradgrad check
self.assertFalse(check_forward_ad, msg="Cannot run forward AD check for gradgradcheck")
for gen_non_contig_grad_outputs in (False, True):
kwargs = {
"gen_non_contig_grad_outputs": gen_non_contig_grad_outputs,
"check_batched_grad": op.check_batched_gradgrad,
"check_grad_dtypes": True,
"nondet_tol": op.gradcheck_nondet_tol,
"fast_mode": op.gradcheck_fast_mode
}
if check == "fwgrad_bwgrad":
kwargs["check_fwd_over_rev"] = True
kwargs["check_rev_over_rev"] = False
kwargs["check_batched_grad"] = False
kwargs["check_undefined_grad"] = False
self.assertTrue(gradgradcheck(fn, gradcheck_args, **kwargs))
else:
self.assertTrue(False, msg="Unknown check requested!")
def _check_helper(self,
device,
dtype,
op,
variant,
check,
*,
check_forward_ad=False,
check_backward_ad=True,
check_batched_grad=None,
check_batched_forward_grad=False):
assert check in ('gradcheck', 'bwgrad_bwgrad', 'fwgrad_bwgrad')
# NB: check_backward_ad does not affect gradgradcheck (always True)
if variant is None:
self.skipTest("Skipped! Variant not implemented.")
if not op.supports_dtype(dtype, torch.device(device).type):
self.skipTest(
f"Skipped! {op.name} does not support dtype {str(dtype)}")
def is_inplace(variant):
if hasattr(variant, "__wrapped__"):
return variant.__wrapped__ is op.get_inplace()
return variant is op.get_inplace()
include_conjugated_inputs = op.test_conjugated_samples and dtype.is_complex
samples = op.sample_inputs(
device,
dtype,
requires_grad=True,
include_conjugated_inputs=include_conjugated_inputs)
for sample in samples:
if sample.broadcasts_input and is_inplace(variant):
continue
# Note on TensorList inputs
#
# gradcheck does not support TensorList inputs so here we pass TensorList
# inputs of size n as n single Tensor inputs to gradcheck and wrap the op
# in a function that puts the n Tensor inputs back into a TensorList
def fn(*inputs):
# Put tensors back into TensorList since we splat them when passing to gradcheck
if is_iterable_of_tensors(sample.input):
n = len(sample.input)
inputs = (inputs[:n], *inputs[n:])
output = op.gradcheck_wrapper(variant, *inputs,
**sample.kwargs)
if sample.output_process_fn_grad is not None:
return sample.output_process_fn_grad(output)
return output
# Splat TensorList inputs into single Tensor inputs
gradcheck_args = (sample.input, ) if isinstance(
sample.input, torch.Tensor) else tuple(sample.input)
gradcheck_args += sample.args
if check == 'gradcheck':
if check_batched_grad is None:
check_batched_grad = op.check_batched_grad
self.assertTrue(
gradcheck(
fn,
gradcheck_args,
check_batched_grad=check_batched_grad,
check_grad_dtypes=True,
nondet_tol=op.gradcheck_nondet_tol,
fast_mode=op.gradcheck_fast_mode,
check_forward_ad=check_forward_ad,
check_backward_ad=check_backward_ad,
check_undefined_grad=True,
check_batched_forward_grad=check_batched_forward_grad))
elif check in ('bwgrad_bwgrad', 'fwgrad_bwgrad'): # gradgrad check
self.assertFalse(
check_forward_ad,
msg="Cannot run forward AD check for gradgradcheck")
for gen_non_contig_grad_outputs in (False, True):
kwargs = {
"gen_non_contig_grad_outputs":
gen_non_contig_grad_outputs,
"check_batched_grad": op.check_batched_gradgrad,
"check_grad_dtypes": True,
"nondet_tol": op.gradcheck_nondet_tol,
"fast_mode": op.gradcheck_fast_mode
}
if check == "fwgrad_bwgrad":
kwargs["check_fwd_over_rev"] = True
kwargs["check_rev_over_rev"] = False
kwargs["check_batched_grad"] = False
kwargs["check_undefined_grad"] = False
self.assertTrue(gradgradcheck(fn, gradcheck_args,
**kwargs))
else:
self.assertTrue(False, msg="Unknown check requested!")
def _test_common(
self,
reduction,
device,
dtype,
unsafe,
axis,
initial_value,
data_arr,
lengths_arr,
expected_arr,
expected_grad_arr,
check_backward,
lengths_dtype=torch.int,
):
lengths = torch.tensor(lengths_arr, device=device, dtype=lengths_dtype)
data = torch.tensor(
data_arr,
device=device,
dtype=dtype,
requires_grad=True,
)
expected_result = torch.tensor(expected_arr,
device=device,
dtype=dtype)
expected_grad = torch.tensor(expected_grad_arr,
device=device,
dtype=dtype)
actual_result = torch.segment_reduce(
data=data,
reduce=reduction,
lengths=lengths,
axis=axis,
unsafe=unsafe,
initial=initial_value,
)
self.assertEqual(expected_result,
actual_result,
rtol=1e-02,
atol=1e-05,
equal_nan=True)
if not check_backward:
return
# Test backward
actual_result.sum().backward()
self.assertEqual(expected_grad,
data.grad,
rtol=1e-02,
atol=1e-05,
equal_nan=True)
# gradcheck does not work well with bfloat16 or fp16 cpu types
# also there is small numerical difference with fp32
if dtype not in [torch.half, torch.bfloat16, torch.float]:
# gradcheck does not like "nan" input, setting to random 10
d_non_nan = np.nan_to_num(data_arr, nan=10)
data = torch.tensor(
# [10 if v == float("nan") else v for v in data],
d_non_nan,
device=device,
dtype=dtype,
requires_grad=True,
)
self.assertTrue(
gradcheck(
lambda x: torch.segment_reduce(
data=x,
reduce=reduction,
lengths=lengths,
axis=axis,
unsafe=unsafe,
initial=initial_value,
),
(data, ),
))
def _test_common(
self,
reduction,
device,
dtype,
unsafe,
axis,
initial_value,
data_arr,
lengths_arr,
expected_arr,
expected_grad_arr,
check_backward,
lengths_dtype=torch.int,
):
lengths = torch.tensor(lengths_arr, device=device, dtype=lengths_dtype)
# generate offsets from lengths
zeros_shape = list(lengths.shape)
zeros_shape[-1] = 1
offsets = torch.cat((lengths.new_zeros(zeros_shape), lengths),
-1).cumsum_(-1)
data = torch.tensor(
data_arr,
device=device,
dtype=dtype,
requires_grad=True,
)
expected_result = torch.tensor(expected_arr,
device=device,
dtype=dtype)
expected_grad = torch.tensor(expected_grad_arr,
device=device,
dtype=dtype)
for mode in ['lengths', 'offsets']:
segment_reduce_kwargs = dict(axis=axis,
unsafe=unsafe,
initial=initial_value)
if (mode == 'lengths'):
segment_reduce_kwargs['lengths'] = lengths
else:
segment_reduce_kwargs['offsets'] = offsets
actual_result = torch.segment_reduce(data=data,
reduce=reduction,
**segment_reduce_kwargs)
self.assertEqual(expected_result,
actual_result,
rtol=1e-02,
atol=1e-05,
equal_nan=True)
if not check_backward:
return
# Test backward
actual_result.sum().backward()
self.assertEqual(expected_grad,
data.grad,
rtol=1e-02,
atol=1e-05,
equal_nan=True)
data = data.clone().detach().requires_grad_(True)
# gradcheck does not work well with bfloat16 or fp16 cpu types
# also there is small numerical difference with fp32
if dtype not in [torch.half, torch.bfloat16, torch.float]:
# gradcheck does not like "nan" input, setting to random 10
d_non_nan = np.nan_to_num(data_arr, nan=10)
new_data = torch.tensor(
# [10 if v == float("nan") else v for v in data],
d_non_nan,
device=device,
dtype=dtype,
requires_grad=True,
)
self.assertTrue(
gradcheck(
lambda x: torch.segment_reduce(
data=x, reduce=reduction, **segment_reduce_kwargs),
(new_data, ),
))
def test_pytorch_scatter_test_cases(self, device, dtypes, reduce):
val_dtype, length_dtype = dtypes
# zero-length segments are filled with reduction inits contrary to pytorch_scatter.
tests = [
{
'src': [1, 2, 3, 4, 5, 6],
'index': [0, 0, 1, 1, 1, 3],
'indptr': [0, 2, 5, 5, 6],
'sum': [3, 12, 0, 6],
'prod': [2, 60, 1, 6],
'mean': [1.5, 4, float('nan'), 6],
'min': [1, 3, float('inf'), 6],
'max': [2, 5, -float('inf'), 6],
},
{
'src': [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]],
'index': [0, 0, 1, 1, 1, 3],
'indptr': [0, 2, 5, 5, 6],
'sum': [[4, 6], [21, 24], [0, 0], [11, 12]],
'prod': [[3, 8], [315, 480], [1, 1], [11, 12]],
'mean': [[2, 3], [7, 8], [float('nan'),
float('nan')], [11, 12]],
'min': [[1, 2], [5, 6], [float('inf'),
float('inf')], [11, 12]],
'max': [[3, 4], [9, 10], [-float('inf'), -float('inf')],
[11, 12]],
},
{
'src': [[1, 3, 5, 7, 9, 11], [2, 4, 6, 8, 10, 12]],
'index': [[0, 0, 1, 1, 1, 3], [0, 0, 0, 1, 1, 2]],
'indptr': [[0, 2, 5, 5, 6], [0, 3, 5, 6, 6]],
'sum': [[4, 21, 0, 11], [12, 18, 12, 0]],
'prod': [[3, 315, 1, 11], [48, 80, 12, 1]],
'mean': [[2, 7, float('nan'), 11], [4, 9, 12,
float('nan')]],
'min': [[1, 5, float('inf'), 11], [2, 8, 12,
float('inf')]],
'max': [[3, 9, -float('inf'), 11], [6, 10, 12, -float('inf')]],
},
{
'src': [[[1, 2], [3, 4], [5, 6]], [[7, 9], [10, 11], [12,
13]]],
'index': [[0, 0, 1], [0, 2, 2]],
'indptr': [[0, 2, 3, 3], [0, 1, 1, 3]],
'sum': [[[4, 6], [5, 6], [0, 0]], [[7, 9], [0, 0], [22, 24]]],
'prod': [[[3, 8], [5, 6], [1, 1]], [[7, 9], [1, 1], [120,
143]]],
'mean': [[[2, 3], [5, 6], [float('nan'),
float('nan')]],
[[7, 9], [float('nan'), float('nan')], [11, 12]]],
'min': [[[1, 2], [5, 6], [float('inf'),
float('inf')]],
[[7, 9], [float('inf'), float('inf')], [10, 11]]],
'max': [[[3, 4], [5, 6], [-float('inf'), -float('inf')]],
[[7, 9], [-float('inf'), -float('inf')], [12, 13]]],
},
{
'src': [[1, 3], [2, 4]],
'index': [[0, 0], [0, 0]],
'indptr': [[0, 2], [0, 2]],
'sum': [[4], [6]],
'prod': [[3], [8]],
'mean': [[2], [3]],
'min': [[1], [2]],
'max': [[3], [4]],
},
{
'src': [[[1, 1], [3, 3]], [[2, 2], [4, 4]]],
'index': [[0, 0], [0, 0]],
'indptr': [[0, 2], [0, 2]],
'sum': [[[4, 4]], [[6, 6]]],
'prod': [[[3, 3]], [[8, 8]]],
'mean': [[[2, 2]], [[3, 3]]],
'min': [[[1, 1]], [[2, 2]]],
'max': [[[3, 3]], [[4, 4]]],
},
]
for test in tests:
data = torch.tensor(test['src'],
dtype=val_dtype,
device=device,
requires_grad=True)
indptr = torch.tensor(test['indptr'],
dtype=length_dtype,
device=device)
dim = indptr.ndim - 1
# calculate lengths from indptr
lengths = torch.diff(indptr, dim=dim)
expected = torch.tensor(test[reduce],
dtype=val_dtype,
device=device)
actual_result = torch.segment_reduce(
data=data,
reduce=reduce,
lengths=lengths,
axis=dim,
unsafe=True,
)
self.assertEqual(actual_result, expected)
# test offsets
actual_result = torch.segment_reduce(
data=data,
reduce=reduce,
offsets=indptr,
axis=dim,
unsafe=True,
)
self.assertEqual(actual_result, expected)
if val_dtype == torch.float64:
def fn(x, mode='lengths'):
initial = 1
# supply initial values to prevent gradcheck from failing for 0 length segments
# where nan/inf are reduction identities that produce nans when calculating the numerical jacobian
if reduce == 'min':
initial = 1000
elif reduce == 'max':
initial = -1000
segment_reduce_args = {x, reduce}
segment_reduce_kwargs = dict(axis=dim,
unsafe=True,
initial=initial)
if mode == 'lengths':
segment_reduce_kwargs[mode] = lengths
elif mode == 'offsets':
segment_reduce_kwargs[mode] = indptr
return torch.segment_reduce(*segment_reduce_args,
**segment_reduce_kwargs)
self.assertTrue(
gradcheck(partial(fn, mode='lengths'),
(data.clone().detach().requires_grad_(True))))
self.assertTrue(
gradcheck(partial(fn, mode='offsets'),
(data.clone().detach().requires_grad_(True))))
def _check_helper(self, device, dtype, op, variant, check):
if variant is None:
self.skipTest("Skipped! Variant not implemented.")
if not op.supports_dtype(dtype, torch.device(device).type):
self.skipTest(
f"Skipped! {op.name} does not support dtype {str(dtype)}")
def is_inplace(variant):
if hasattr(variant, "__wrapped__"):
return variant.__wrapped__ is op.get_inplace()
return variant is op.get_inplace()
samples = op.sample_inputs(device, dtype, requires_grad=True)
for sample in samples:
if sample.broadcasts_input and is_inplace(variant):
continue
# Note on TensorList inputs
#
# gradcheck does not support TensorList inputs so here we pass TensorList
# inputs of size n as n single Tensor inputs to gradcheck and wrap the op
# in a function that puts the n Tensor inputs back into a TensorList
def fn(*inputs):
# Put tensors back into TensorList since we splat them when passing to gradcheck
if is_iterable_of_tensors(sample.input):
n = len(sample.input)
inputs = (inputs[:n], *inputs[n:])
output = op.gradcheck_wrapper(variant, *inputs,
**sample.kwargs)
if sample.output_process_fn_grad is not None:
return sample.output_process_fn_grad(output)
return output
# Splat TensorList inputs into single Tensor inputs
gradcheck_args = (sample.input, ) if isinstance(
sample.input, torch.Tensor) else tuple(sample.input)
gradcheck_args += sample.args
if check == 'gradcheck':
self.assertTrue(
gradcheck(fn,
gradcheck_args,
check_batched_grad=op.check_batched_grad,
check_grad_dtypes=True,
nondet_tol=op.gradcheck_nondet_tol,
fast_mode=op.gradcheck_fast_mode))
elif check == 'gradgradcheck':
self.assertTrue(
gradgradcheck(fn,
gradcheck_args,
gen_non_contig_grad_outputs=False,
check_batched_grad=op.check_batched_gradgrad,
check_grad_dtypes=True,
nondet_tol=op.gradcheck_nondet_tol,
fast_mode=op.gradcheck_fast_mode))
self.assertTrue(
gradgradcheck(fn,
gradcheck_args,
gen_non_contig_grad_outputs=True,
check_batched_grad=op.check_batched_gradgrad,
check_grad_dtypes=True,
nondet_tol=op.gradcheck_nondet_tol,
fast_mode=op.gradcheck_fast_mode))
else:
self.assertTrue(False, msg="Unknown check requested!")
def _test_simple_1d(self, reduction, device, dtype, unsafe, axis):
lengths = torch.tensor([1, 2, 3, 0], device=device)
data = torch.tensor(
[1, float("nan"), 3, 4, 5, 5],
device=device,
dtype=dtype,
requires_grad=True,
)
initial_value = 0
if reduction == "max":
expected_result = torch.tensor(
[1, float("nan"), 5, initial_value], device=device, dtype=dtype
)
expected_grad = torch.tensor(
[1, 1, 0, 0, 0.5, 0.5], device=device, dtype=dtype
)
elif reduction == "mean":
expected_result = torch.tensor(
[1, float("nan"), 4.666, initial_value], device=device, dtype=dtype
)
expected_grad = torch.tensor(
[1.0, 0.5, 0.5, 0.333, 0.333, 0.333], device=device, dtype=dtype
)
actual_result = torch.segment_reduce(
data=data,
reduce=reduction,
lengths=lengths,
axis=axis,
unsafe=unsafe,
initial=initial_value,
)
self.assertEqual(
expected_result, actual_result, rtol=1e-02, atol=1e-05, equal_nan=True
)
# TODO: Remove this check once cuda backward support is implemented
if data.is_cuda:
return
# Test backward
actual_result.sum().backward()
self.assertEqual(
expected_grad, data.grad, rtol=1e-02, atol=1e-05, equal_nan=True
)
# gradcheck does not work well with bfloat16 or fp16 cpu types
# also there is small numerical difference with fp32
if dtype not in [torch.half, torch.bfloat16, torch.float]:
# gradcheck does not like "nan" input
data = torch.tensor(
[1, 10, 3, 4, 5, 5],
device=device,
dtype=dtype,
requires_grad=True,
)
self.assertTrue(
gradcheck(
lambda x: torch.segment_reduce(
data=x,
reduce=reduction,
lengths=lengths,
axis=axis,
unsafe=unsafe,
initial=initial_value,
),
(data,),
)
)
def _test_max_simple_1d(self, device, dtype, unsafe, axis):
lengths = torch.tensor([1, 2, 3, 0], device=device)
data = torch.tensor(
[1, float("nan"), 3, 4, 5, 5],
device=device,
dtype=dtype,
requires_grad=True,
)
initial_value = 0
expected_result = torch.tensor([1, float("nan"), 5, initial_value],
device=device,
dtype=dtype)
actual_result = torch.segment_reduce(
data=data,
reduce="max",
lengths=lengths,
axis=axis,
unsafe=unsafe,
initial=initial_value,
)
self.assertEqual(expected_result,
actual_result,
rtol=1e-03,
atol=1e-05,
equal_nan=True)
# Backward is only supported for cpu tensors for now. Return early if cuda
if data.is_cuda:
return
# Test backward
expected_grad = torch.tensor([1, 1, 0, 0, 0.5, 0.5],
device=device,
dtype=dtype)
actual_result.sum().backward()
self.assertEqual(expected_grad,
data.grad,
rtol=1e-03,
atol=1e-05,
equal_nan=True)
# gradcheck does not work well with bfloat16 or fp16 cpu types
# also there is small numerical difference with fp32
if dtype not in [torch.half, torch.bfloat16, torch.float]:
# gradcheck does not like "nan" input
data = torch.tensor(
[1, 10, 3, 4, 5, 5],
device=device,
dtype=dtype,
requires_grad=True,
)
self.assertTrue(
gradcheck(
lambda x: torch.segment_reduce(
data=x,
reduce="max",
lengths=lengths,
axis=axis,
unsafe=unsafe,
initial=initial_value,
),
(data, ),
))
