Add ability to check shapes with wildcards

[KristianHolsheimer]

I often find myself writing the following sort of thing:

chex.assert_rank(x, 2)
x.shape[1] == num_actions, "some custom message ..."

It would be nice to be able to simply check the shape with a wildcard, i.e.

chex.assert_shape(x, (None, num_actions))

What do you think?

[KristianHolsheimer]

Also, I tried to run the tests locally with

pytest .

which results in failures saying Trying to access flag --chex_n_cpu_devices before flags were parsed.

What do you use to run your tests?

[hbq1]

Hey Kristian, this is a great idea 👍 Thanks a lot for the contribution!

You can run tests as simple as python <tests_src.py>. To run all tests at once you can execute e.g.

for test in `ls chex/_src/*_test.py`; do if ! `python $test`; then break; fi; done

if you work in bash.
That to be said, some tests may require additional dependencies, e.g. tree or dm-tree.