question-mark
Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

[Feature suggestion] Add a simple shape assertion in einops-style notation

See original GitHub issue

I find that when I use einops, I end up with a mix of einops operations and regular asserts like this:

x = repeat(target, "b t h w c -> b t k h w c", k=K)
x = do_thing(x)
assert x.shape == (batch_size, T, H, W, C)

I love the einops notation, and I can remove shape assertions immediately before or after an einops operation. But when I don’t need to do a reshape/repeat/etc, I have to fall back to the assert notation to check the shape. I like to include lots of shape asserts in general both to make sure I haven’t accidentally included a bug, but also for improving readability, so the reader always knows the shape of tensors. Asserts are superior to comments, as they will fail if you forget to update them, ensuring that they’re always accurate.

So I propose a new einops “operation”, which does nothing except check shapes, and would raise an assertion error if the shape is incorrect. It would have a notation analogous to other einops operations:

from einops import check_shape
check_shape(x, "b t h w c", t=T, h=H, c=3)

This is preferable to the normal assert for a few reasons:

  • if we’re already using einops, it’s nice to have a standard notation format, rather than mixing two notation formats. It makes the code more readable.
  • i like the “b t h w c” style notation better than the assert-style notation, it allows you to give a “name” to each axis as opposed to just specifying its value.
  • this notation allows you to only check certain axes. eg I don’t normally care to check the batch_size dim, but doing assert x.shape[1:] = (T, H, W, C) is kinda yucky - and worse for axes not at the beginning or end.

Issue Analytics

  • State:open
  • Created 2 years ago
  • Comments:5 (1 by maintainers)

github_iconTop GitHub Comments

4reactions
nils-wernercommented, Feb 17, 2022

I didn’t check the performance implications, but for now you may be able to get away with

def check_shape(tensor, pattern, **kwargs):
    return einops.rearrange(tensor, f"{pattern} -> {pattern}", **kwargs)
2reactions
sebastianjaszczurcommented, Feb 7, 2022

I’ve personally found useful assert_shape as a decorator for a function/layer. Basically something like

@assert_shape('...a->...b')
class Dense(nn.Layer):
    (...)

The above means that Dense should have the same signature as einsum(‘…a->…b’). It should check that the last dimension is indeed of the same size etc.

I’ve implemented this idea in Trax, here (there are aksi sine examples and documentation): https://github.com/google/trax/blob/master/trax/layers/assert_shape.py This is certainly implementable also in PyTorch; I’m not sure about other frameworks.

Read more comments on GitHub >

github_iconTop Results From Across the Web

What does assert equal to empty parentheses (assert x ...
() is an empty tuple, this assert checks if the shape of cost is an empty tuple. So it checks if cost is...
Read more >

github_iconTop Related Medium Post

No results found

github_iconTop Related StackOverflow Question

No results found

github_iconTroubleshoot Live Code

Lightrun enables developers to add logs, metrics and snapshots to live code - no restarts or redeploys required.
Start Free

github_iconTop Related Reddit Thread

No results found

github_iconTop Related Hackernoon Post

No results found

github_iconTop Related Tweet

No results found

github_iconTop Related Dev.to Post

No results found

github_iconTop Related Hashnode Post

No results found