einops
einops copied to clipboard
[Extend feature] Reduce and repeat not supporting tensor concatenation.
Einops allows the use of rearrange to concatenate a list of tensors, like shown in the einops for deep learning notebook.
import torch
import einops
x = [torch.rand([3,20,20]) for _ in range(10)]
stacked_and_reshaped = einops.rearrange(x, "b c h w -> b c (h w)")
However, this functionality does not work for reduce and repeat:
stacked_and_averaged = einops.reduce(x, "b c h w -> c h w", reduction="mean")
stacked_and_repeated = einops.repeat(x, "b c h w -> why_not b c h w", why_not=2)
both operations produce these error messages (this is obviously for repeat)
File "/home/myName/miniconda3/envs/myEnv/lib/python3.9/site-packages/einops/einops.py", line 536, in repeat
return reduce(tensor, pattern, reduction='repeat', **axes_lengths)
File "/home/myName/miniconda3/envs/myEnv/lib/python3.9/site-packages/einops/einops.py", line 409, in reduce
return _apply_recipe(recipe, tensor, reduction_type=reduction)
File "/home/myName/miniconda3/envs/myEnv/lib/python3.9/site-packages/einops/einops.py", line 230, in _apply_recipe
backend = get_backend(tensor)
File "/home/myName/miniconda3/envs/myEnv/lib/python3.9/site-packages/einops/_backends.py", line 52, in get_backend
raise RuntimeError('Tensor type unknown to einops {}'.format(type(tensor)))
RuntimeError: Tensor type unknown to einops <class 'list'>
The expected behavior is of course, that it would work like with rearrange.
Your platform Python: 3.9 einops: 0.4.0 DL_package: pyTorch
see #172 , that's a plan for the next release