def _test_property(self, fn):
for constructor in _iter_constructors():
num_nested_tensor = 3
nested_tensor_lists = [utils.gen_nested_list(i, i, 3)
for i in range(1, num_nested_tensor)]
first_tensors = [utils.get_first_tensor(
ntl) for ntl in nested_tensor_lists]
nested_tensors = [constructor(ntl) for ntl in nested_tensor_lists]
for nested_tensor, first_tensor in zip(nested_tensors, first_tensors):
self.assertEqual(fn(nested_tensor), fn(first_tensor))
def _test_property(self, fn):
for constructor in _iter_constructors():
# TODO: Used to be 3. Currently only supporting nested dim 1.
num_nested_tensor = 1
nested_tensor_lists = [utils.gen_nested_list(i, i, 3)
for i in range(1, num_nested_tensor)]
first_tensors = [utils.get_first_tensor(
ntl) for ntl in nested_tensor_lists]
nested_tensors = [constructor(ntl) for ntl in nested_tensor_lists]
for nested_tensor, first_tensor in zip(nested_tensors, first_tensors):
self.assertEqual(fn(nested_tensor), fn(first_tensor))
def _test_unary(self):
data = utils.gen_nested_list(1, nested_dim, 3)
data = utils.nested_map(lambda x: x.to(device), data)
if func__ in ['log', 'log10', 'log2', 'rsqrt', 'sqrt']:
data = utils.nested_map(lambda x: x.abs(), data)
if func__ in ['acos', 'asin', 'erfinv', 'log1p']:
data = utils.nested_map(lambda x: x.clamp(min=0, max=1), data)
if func__ in ['mvlgamma']:
data = utils.nested_map(lambda x: x.clamp(min=1), data)
func_ = getattr(torch, func__)
method_ = getattr(nestedtensor.NestedTensor, func__)
method_inplace_ = getattr(nestedtensor.NestedTensor, func__ + "_")
if func__ in ['clamp']:
def func(x, out=None):
return func_(x, min=-1, max=1, out=out)
def method(x):
return method_(x, min=-1, max=1)
def method_inplace(x):
return method_inplace_(x, min=-1, max=1)
elif func__ in ['clamp_min']:
def func(x, out=None):
return func_(x, min=-1, out=out)
def method(x):
return method_(x, min=-1)
def method_inplace(x):
return method_inplace_(x, min=-1)
elif func__ in ['clamp_max']:
def func(x, out=None):
return func_(x, 1, out=out)
def method(x):
return method_(x, 1)
def method_inplace(x):
return method_inplace_(x, 1)
elif func__ in ['mvlgamma']:
def func(x):
return func_(x, p=2)
def method(x):
return method_(x, p=2)
def method_inplace(x):
return method_inplace_(x, p=2)
elif func__ in ['renorm']:
def func(x, out=None):
return func_(x, 2, 0, 1.0, out=out)
def method(x):
return method_(x, 2, 0, 1.0)
def method_inplace(x):
return method_inplace_(x, 2, 0, 1.0)
elif func__ in ['fmod']:
def func(x, out=None):
return func_(x, 0.3, out=out)
def method(x):
return method_(x, 0.3)
def method_inplace(x):
return method_inplace_(x, 0.3)
else:
func = func_
method = method_
method_inplace = method_inplace_
def _close(t1, t2):
self.assertAlmostEqual(t1, t2, ignore_contiguity=True)
a1 = ntnt(data, device=device)
a2 = ntnt(utils.nested_map(func, data), device=device)
_close(func(a1), a2)
_close(method(a1), a2)
a1 = ntnt_nograd(data, device=device)
a2 = ntnt_nograd(utils.nested_map(func, data), device=device)
a3 = ntnt_nograd(data, device=device)
self.assertEqual(a1.nested_dim(), a2.nested_dim())
self.assertEqual(a2.nested_dim(), a3.nested_dim())
if func__ not in ['mvlgamma']:
func(a1, out=a3)
# TODO: Abstract this
_close(func(a1), a3)
_close(method_inplace(a1), a2)
_close(a1, a2)
