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)