chex.Dimensions API enhancement

I would like to propose an API enhancement that allow the use of chex.Dimensions inside function annotations. If there is interest I'd like to contribute. Example below:
```
dims = chex.Dimensions(B=batch_size, T=sequence_len, E=embedding_dim)
...
def foo(arr: chex.Array):
chex.assert_shape(arr, dims['BTE'])
# fn logic

### turns into ###

def foo(arr: chex.Array(dims['BTE'])): # behind the scenes assert on function call
# fn logic

```
This is particularly useful for dataclasses e.g.
```
dims = chex.Dimensions(B=batch_size, T=rollout_len)

# asserts are run on instantiation
class TimeStep:
q_values: chex.Array(dims['BT'])
discounts: chex.Array(dims['BT'])
rewards: chex.Array(dims['BT'])
```
Pros:
- reduces clutter that asserts can add
- allows user to view the shape expected by function or class in editor (not sure what you call the VScode popup)
- example: using RLax, in order to know what shape is expected for each arg in a loss fn you need to either look at source code or wait for fn call to raise an assert

Cons:
- increased API complexity
- ...?