def test_parse_shape_symbolic():
backends = collect_test_backends(symbolic=True, layers=False)
backends += collect_test_backends(symbolic=True, layers=True)
for backend in backends:
if backend.framework_name == 'keras':
# need special way to compile, shape vars can be used only inside layers
continue
print('special shape parsing for', backend.framework_name)
input_symbols = [
backend.create_symbol([10, 20, 30, 40]),
backend.create_symbol([10, 20, None, None]),
backend.create_symbol([None, None, None, None]),
]
if backend.framework_name in ['mxnet.symbol']:
# mxnet can't normally run inference
input_symbols = [backend.create_symbol([10, 20, 30, 40])]
for input_symbol in input_symbols:
shape_placeholder = parse_shape(input_symbol, 'a b c d')
shape = {}
for name, symbol in shape_placeholder.items():
shape[name] = symbol if isinstance(symbol, int) \
else backend.eval_symbol(symbol, [(input_symbol, numpy.zeros([10, 20, 30, 40]))])
print(shape)
result_placeholder = rearrange(
input_symbol,
'a b (c1 c2) (d1 d2) -> (a b d1) c1 (c2 d2)',
**parse_shape(input_symbol, 'a b c1 _'),
d2=2)
result = backend.eval_symbol(
result_placeholder,
[(input_symbol, numpy.zeros([10, 20, 30, 40]))])
print(result.shape)
assert result.shape == (10 * 20 * 20, 30, 1 * 2)
assert numpy.allclose(result, 0)
def test6(x):
# parsing parameters
t = rearrange(x, 'b c h w -> (b h w) c')
t = t[:, ::
2] # replacement for dot-product, just changes size of second axis
assert t.shape == (10 * 30 * 40, 10)
y = rearrange(t, '(b h w) c2 -> b c2 h w', **parse_shape(x, 'b _ h w'))
assert y.shape == (10, 10, 30, 40)
return y
def test_parse_shape_symbolic(self, shape):
print('special shape parsing for', self.backend.framework_name)
if self.backend.framework_name in ['mxnet.symbol']:
# mxnet can't normally run inference
shape = [10, 20, 30, 40]
input_symbol = self.backend.create_symbol(shape)
shape_placeholder = parse_shape(input_symbol, 'a b c d')
shape = {}
for name, symbol in shape_placeholder.items():
shape[name] = symbol if isinstance(symbol, int) \
else self.backend.eval_symbol(symbol, [(input_symbol, numpy.zeros([10, 20, 30, 40]))])
print(shape)
result_placeholder = rearrange(input_symbol, 'a b (c1 c2) (d1 d2) -> (a b d1) c1 (c2 d2)',
**parse_shape(input_symbol, 'a b c1 _'), d2=2)
result = self.backend.eval_symbol(result_placeholder, [(input_symbol, numpy.zeros([10, 20, 30, 40]))])
print(result.shape)
assert result.shape == (10 * 20 * 20, 30, 1 * 2)
assert numpy.allclose(result, 0)
def test_ellipsis(self, static_shape: List[int], shape: List[Optional[int]],
pattern: str, expected: Dict[str, int]):
if self.backend.framework_name in ['mxnet.symbol']:
# mxnet can't normally run inference
shape = static_shape
input_symbol = self.backend.create_symbol(shape)
shape_placeholder = parse_shape(input_symbol, pattern)
out_shape = {}
for name, symbol in shape_placeholder.items():
if isinstance(symbol, int):
out_shape[name] = symbol
else:
out_shape[name] = self.backend.eval_symbol(
symbol, [(input_symbol, numpy.zeros(static_shape))])
assert out_shape == expected
def test_parse_shape_imperative():
backends = collect_test_backends(symbolic=False, layers=False)
backends += collect_test_backends(symbolic=False, layers=True)
for backend in backends:
print('Shape parsing for ', backend.framework_name)
x = numpy.zeros([10, 20, 30, 40])
parsed1 = parse_shape(x, 'a b c d')
parsed2 = parse_shape(backend.from_numpy(x), 'a b c d')
assert parsed1 == parsed2 == dict(a=10, b=20, c=30, d=40)
assert parsed1 != dict(a=1, b=20, c=30, d=40) != parsed2
parsed1 = parse_shape(x, '_ _ _ _')
parsed2 = parse_shape(backend.from_numpy(x), '_ _ _ _')
assert parsed1 == parsed2 == dict()
parsed1 = parse_shape(x, '_ _ _ hello')
parsed2 = parse_shape(backend.from_numpy(x), '_ _ _ hello')
assert parsed1 == parsed2 == dict(hello=40)
parsed1 = parse_shape(x, '_ _ a1 a1a111a')
parsed2 = parse_shape(backend.from_numpy(x), '_ _ a1 a1a111a')
assert parsed1 == parsed2 == dict(a1=30, a1a111a=40)
def test_underscore_one(self):
parsed1 = parse_shape(self.x, '_ _ _ hello')
parsed2 = parse_shape(self.backend.from_numpy(self.x), '_ _ _ hello')
assert parsed1 == parsed2 == dict(hello=40)
def test_parse_shape_imperative(self):
print('Shape parsing for ', self.backend.framework_name)
parsed1 = parse_shape(self.x, 'a b c d')
parsed2 = parse_shape(self.backend.from_numpy(self.x), 'a b c d')
assert parsed1 == parsed2 == dict(a=10, b=20, c=30, d=40)
assert parsed1 != dict(a=1, b=20, c=30, d=40) != parsed2
def test_ellipsis(self, shape: List[int], pattern: str,
expected: Dict[str, int]):
x = numpy.ones(shape)
parsed1 = parse_shape(x, pattern)
parsed2 = parse_shape(self.backend.from_numpy(x), pattern)
assert parsed1 == parsed2 == expected
def test_repeating(self):
with assert_raises(einops.EinopsError):
parse_shape(self.x, 'a a b b')
with assert_raises(einops.EinopsError):
parse_shape(self.backend.from_numpy(self.x), 'a a b b')
def test_underscore_several(self):
parsed1 = parse_shape(self.x, '_ _ a1 a1a111a')
parsed2 = parse_shape(self.backend.from_numpy(self.x), '_ _ a1 a1a111a')
assert parsed1 == parsed2 == dict(a1=30, a1a111a=40)
