def test_gradients_imperatives():
# lazy - just checking reductions
for reduction in _reductions:
x = numpy.arange(1, 1 + 2 * 3 * 4).reshape([2, 3, 4]).astype('float32')
results = {}
for backend in imp_op_backends:
y0 = backend.from_numpy(x)
if not hasattr(y0, 'grad'):
continue
if 'mxnet' in backend.framework_name:
backend.mx.autograd.set_recording(True)
y1 = reduce(y0, 'a b c -> c a', reduction=reduction)
y2 = reduce(y1, 'c a -> a c', reduction=reduction)
y3 = reduce(y2, 'a (c1 c2) -> a', reduction=reduction, c1=2)
y4 = reduce(y3, '... -> ', reduction=reduction)
if 'mxnet' in backend.framework_name:
backend.mx.autograd.set_recording(False)
y4.backward()
grad = backend.to_numpy(y0.grad)
results[backend.framework_name] = grad
print('comparing gradients for', results.keys())
for name1, grad1 in results.items():
for name2, grad2 in results.items():
assert numpy.allclose(grad1, grad2), [
name1, name2, 'provided different gradients'
]
def check_reversion(x, repeat_pattern, **sizes):
"""Checks repeat pattern by running reduction """
left, right = repeat_pattern.split('->')
reduce_pattern = right + '->' + left
repeated = reduce(x, repeat_pattern, reduction='repeat', **sizes)
reduced_min = reduce(repeated, reduce_pattern, reduction='min', **sizes)
reduced_max = reduce(repeated, reduce_pattern, reduction='max', **sizes)
assert numpy.array_equal(x, reduced_min)
assert numpy.array_equal(x, reduced_max)
def test_reduction_with_callable_imperatives():
x_numpy = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5,
6]).astype('float32')
x_numpy /= x_numpy.max()
def logsumexp_torch(x, tuple_of_axes):
return x.logsumexp(tuple_of_axes)
def logsumexp_tf(x, tuple_of_axes):
import tensorflow as tf
return tf.reduce_logsumexp(x, tuple_of_axes)
def logsumexp_chainer(x, tuple_of_axes):
import chainer
return chainer.functions.logsumexp(x, tuple_of_axes)
def logsumexp_keras(x, tuple_of_axes):
import tensorflow.keras.backend as k
return k.logsumexp(x, tuple_of_axes)
def logsumexp_numpy(x, tuple_of_axes):
# very naive logsumexp to compare to
minused = x.max(tuple_of_axes)
y = x - x.max(tuple_of_axes, keepdims=True)
y = numpy.exp(y)
y = numpy.sum(y, axis=tuple_of_axes)
return numpy.log(y) + minused
from einops._backends import TorchBackend, ChainerBackend, TensorflowBackend, KerasBackend, NumpyBackend
backend2callback = {
TorchBackend.framework_name: logsumexp_torch,
ChainerBackend.framework_name: logsumexp_chainer,
TensorflowBackend.framework_name: logsumexp_tf,
KerasBackend.framework_name: logsumexp_keras,
NumpyBackend.framework_name: logsumexp_numpy,
}
for backend in imp_op_backends:
if backend.framework_name not in backend2callback:
continue
backend_callback = backend2callback[backend.framework_name]
x_backend = backend.from_numpy(x_numpy)
for pattern1, pattern2 in equivalent_reduction_patterns:
print('Test reduction with callable for ', backend.framework_name,
pattern1, pattern2)
output_numpy = reduce(x_numpy, pattern1, reduction=logsumexp_numpy)
output_backend = reduce(x_backend,
pattern1,
reduction=backend_callback)
assert numpy.allclose(
output_numpy,
backend.to_numpy(output_backend),
)
def test_repeat_numpy():
# check repeat vs reduce. Repeat works ok if reverse reduction with min and max work well
x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5])
x1 = reduce(x, 'a b c -> copy a b c ', reduction='repeat', copy=1)
assert numpy.array_equal(x[None], x1)
for pattern, axis_dimensions in repeat_test_cases:
check_reversion(x, pattern, **axis_dimensions)
def test9(x):
# squeeze - unsqueeze
y = reduce(x, 'b c h w -> b c () ()', reduction='max')
assert y.shape == (10, 20, 1, 1)
y = rearrange(y, 'b c () () -> c b')
assert y.shape == (20, 10)
return y
def test_repeat_symbolic():
x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5])
for backend in sym_op_backends:
print('Repeat tests for ', backend.framework_name)
for pattern, axis_dimensions in repeat_test_cases:
expected = reduce(x,
pattern,
reduction='repeat',
**axis_dimensions)
sym = backend.create_symbol(x.shape)
result = backend.eval_symbol(
reduce(sym, pattern, reduction='repeat', **axis_dimensions),
[[sym, x]])
assert numpy.array_equal(result, expected)
def test_repeat_imperatives():
x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5])
for backend in imp_op_backends:
print('Repeat tests for ', backend.framework_name)
for pattern, axis_dimensions in repeat_test_cases:
expected = reduce(x,
pattern,
reduction='repeat',
**axis_dimensions)
converted = backend.from_numpy(x)
repeated = reduce(converted,
pattern,
reduction='repeat',
**axis_dimensions)
result = backend.to_numpy(repeated)
assert numpy.array_equal(result, expected)
def test_ellipsis_ops_numpy():
x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
for pattern in identity_patterns:
assert numpy.array_equal(x, rearrange(x, pattern)), pattern
for pattern1, pattern2 in equivalent_rearrange_patterns:
assert numpy.array_equal(rearrange(x, pattern1),
rearrange(x, pattern2))
for reduction in ['min', 'max', 'sum']:
for pattern1, pattern2 in equivalent_reduction_patterns:
assert numpy.array_equal(reduce(x, pattern1, reduction=reduction),
reduce(x, pattern2, reduction=reduction))
# now just check coincidence with numpy
all_rearrange_patterns = [*identity_patterns]
for pattern_pairs in equivalent_rearrange_patterns:
all_rearrange_patterns.extend(pattern_pairs)
def test8(x):
# max-pooling
y = reduce(x,
'b c (h h1) (w w1) -> b c h w',
reduction='max',
h1=2,
w1=2)
assert y.shape == (10, 20, 30 // 2, 40 // 2)
return y
def test_reduction_imperatives():
for backend in imp_op_backends:
print('Reduction tests for ', backend.framework_name)
for reduction in _reductions:
# slight redundancy for simpler order - numpy version is evaluated multiple times
input = numpy.arange(2 * 3 * 4 * 5 * 6,
dtype='int64').reshape([2, 3, 4, 5, 6])
if reduction in ['mean', 'prod']:
input = input / input.astype('float64').mean()
test_cases = [
['a b c d e -> ', {},
getattr(input, reduction)()],
['a ... -> ', {}, getattr(input, reduction)()],
[
'(a1 a2) ... (e1 e2) -> ',
dict(a1=1, e2=2),
getattr(input, reduction)()
],
[
'a b c d e -> (e c) a', {},
getattr(input,
reduction)(axis=(1,
3)).transpose(2, 1,
0).reshape([-1, 2])
],
[
'a ... c d e -> (e c) a', {},
getattr(input,
reduction)(axis=(1,
3)).transpose(2, 1,
0).reshape([-1, 2])
],
[
'a b c d e ... -> (e c) a', {},
getattr(input,
reduction)(axis=(1,
3)).transpose(2, 1,
0).reshape([-1, 2])
],
[
'a b c d e -> (e c a)', {},
getattr(input,
reduction)(axis=(1, 3)).transpose(2, 1,
0).reshape([-1])
],
['(a a2) ... -> (a2 a) ...',
dict(a2=1), input],
]
for pattern, axes_lengths, expected_result in test_cases:
result = reduce(backend.from_numpy(input.copy()),
pattern,
reduction=reduction,
**axes_lengths)
result = backend.to_numpy(result)
assert numpy.allclose(result, expected_result)
def test_repeat_numpy():
x = numpy.arange(2 * 3 * 5).reshape(2, 3, 5)
x1 = reduce(x, 'a b c -> copy a b c ', reduction='repeat', copy=1)
assert numpy.array_equal(x[None], x1)
def check_reversion(x, repeat_pattern, **sizes):
left, right = repeat_pattern.split('->')
reduce_pattern = right + '->' + left
repeated = reduce(x, repeat_pattern, reduction='repeat', **sizes)
reduced_min = reduce(repeated,
reduce_pattern,
reduction='min',
**sizes)
reduced_max = reduce(repeated,
reduce_pattern,
reduction='max',
**sizes)
assert numpy.array_equal(x, reduced_min)
assert numpy.array_equal(x, reduced_max)
for pattern, axis_dimensions in repeat_test_cases:
check_reversion(x, pattern, **axis_dimensions)
def test_reduction_stress_imperatives():
for backend in imp_op_backends:
print('Stress-testing reduction for ', backend.framework_name)
for reduction in _reductions + ('rearrange', ):
dtype = 'int64'
coincide = numpy.array_equal
if reduction in ['mean', 'prod']:
dtype = 'float64'
coincide = numpy.allclose
for n_axes in range(6 if 'mxnet' in backend.framework_name else (
7 if 'oneflow' in backend.framework_name else 11)):
shape = numpy.random.randint(2, 4, size=n_axes)
permutation = numpy.random.permutation(n_axes)
skipped = 0 if reduction == 'rearrange' else numpy.random.randint(
n_axes + 1)
left = ' '.join('x' + str(i) for i in range(n_axes))
right = ' '.join('x' + str(i) for i in permutation[skipped:])
pattern = left + '->' + right
x = numpy.arange(1, 1 + numpy.prod(shape),
dtype=dtype).reshape(shape)
if reduction == 'prod':
x /= x.mean() # to avoid overflows
result1 = reduce(x, pattern, reduction=reduction)
result2 = x.transpose(permutation)
if skipped > 0:
result2 = getattr(result2,
reduction)(axis=tuple(range(skipped)))
assert coincide(result1, result2)
if n_axes == 0 and 'mxnet' in backend.framework_name:
# known mxnet bug, can't attach gradients to scalar
continue
check_op_against_numpy(backend,
x,
pattern,
reduction=reduction,
axes_lengths={},
is_symbolic=False)
def operation(x):
if reduction == 'rearrange':
return rearrange(x, pattern, **axes_lengths)
else:
return reduce(x, pattern, reduction, **axes_lengths)
def test_reduction_symbolic():
for backend in sym_op_backends:
print('Reduction tests for ', backend.framework_name)
for reduction in _reductions:
input = numpy.arange(2 * 3 * 4 * 5 * 6,
dtype='int64').reshape([2, 3, 4, 5, 6])
input = input / input.astype('float64').mean()
# slight redundancy for simpler order - numpy version is evaluated multiple times
test_cases = [
['a b c d e -> ', {},
getattr(input, reduction)()],
['a ... -> ', {}, getattr(input, reduction)()],
[
'(a a2) ... (e e2) -> ',
dict(a2=1, e2=1),
getattr(input, reduction)()
],
[
'a b c d e -> (e c) a', {},
getattr(input,
reduction)(axis=(1,
3)).transpose(2, 1,
0).reshape([-1, 2])
],
[
'a ... c d e -> (e c) a', {},
getattr(input,
reduction)(axis=(1,
3)).transpose(2, 1,
0).reshape([-1, 2])
],
[
'a b c d e ... -> (e c) a', {},
getattr(input,
reduction)(axis=(1,
3)).transpose(2, 1,
0).reshape([-1, 2])
],
[
'a b c d e -> (e c a)', {},
getattr(input,
reduction)(axis=(1, 3)).transpose(2, 1,
0).reshape([-1])
],
['(a a2) ... -> (a2 a) ...',
dict(a2=1), input],
]
for pattern, axes_lengths, expected_numpy_result in test_cases:
shapes = [input.shape]
if backend.framework_name != 'mxnet.symbol':
# mxnet can't handle non-specified shapes
shapes.append([None for _ in input.shape])
for shape in shapes:
sym = backend.create_symbol(shape)
result_sym = reduce(sym,
pattern,
reduction=reduction,
**axes_lengths)
result = backend.eval_symbol(result_sym, [(sym, input)])
assert numpy.allclose(result, expected_numpy_result)
if True:
shape = []
_axes_lengths = {**axes_lengths}
for axis, length in zip('abcde', input.shape):
# filling as much as possible with Nones
if axis in pattern:
shape.append(None)
_axes_lengths[axis] = length
else:
shape.append(length)
sym = backend.create_symbol(shape)
result_sym = reduce(sym,
pattern,
reduction=reduction,
**_axes_lengths)
result = backend.eval_symbol(result_sym, [(sym, input)])
assert numpy.allclose(result, expected_numpy_result)
