Allow access to local variables in shape constraints

[jesnie]

Feature request

Allow access to local variables in shape constraints.

Motivation

For example, to be able to define that shapes of parameters must match class variables:

class MyModel:
    def __init__(self, ...):
        self.n_features = ...
        self.n_outputs = ...

    @check_shapes(
        "train_features: [n_rows, $self.n_features]",
        "train_labels: [n_rows, $self.n_outputs]",
    )
    def train(self, train_features, train_labels):
        ...

Proposal

Some way of being able to refer to local variables. We'd need to be able to distinguish local variables from the usual shapes variables. In the above example I've used $, but we could use some other syntax. Or maybe we can magically guess?

Maybe we can get inspiration for syntax from the popular einops library: https://github.com/arogozhnikov/einops