assert_tree_shape_prefix requires tuple instead of Sequence

[hylkedonker]

Hi,

Thanks for sharing this nice package! I especially like all the assertions for pytrees.
However, I came across the following inconsistency in the documentation:
Current situation
According to the docs, the shape_prefix argument of assert_tree_shape_prefix is of type Sequence[int].
However, when I pass sequence such as list (instead of a tuple)

import chex
import jax.numpy as jnp

mytree = {'a': jnp.array([[[1], [2]]])}
chex.assert_tree_shape_prefix(mytree, shape_prefix=[1, 2])  # AssertionError!

The assertion raises an exception:

AssertionError: [Chex] Assertion assert_tree_shape_prefix failed: Tree leaf 'a' has a shape prefix different from expected: (1, 2) != [1, 2].

The error can simply be fixed by using a tuple instead:

chex.assert_tree_shape_prefix(mytree, shape_prefix=(1, 2))  # OK!

I think this is not the only inconsistent function, but I did not check for others.

Desired situation
Ideally, I would like the function to behave like in the docs, so that I can also pass a list. Why? I think being able to choose square brackets (i.e., list) after closing parenthesis helps readability.

I would be interested to hear your opinion and I am happy to contribute a pull request.

Keep up the good work!

Hylke

[stompchicken]

Hi Hylke, thanks for raising the issue! I agree with you that assert_tree_shape_prefix ought to work with sequences, as it claims to in the type annotation. If you would like to contribute a pull request, that would be great.

I think the best thing to do would be to convert the shape_prefix argument to a tuple inside assert_tree_shape_prefix before the assert_fn is defined. Ideally you would also add a unit test :)

[hylkedonker]

Thanks for your prompt response.
I'Il try to give it a go the coming days. :-)