Explicit way of register classes

Currently, `chex.assert_xyz` uses `tree` under the hood. This make it impossible to use it with Jax trees (`jax.tree_utils.register`) which do not support `collections.abc.Mapping`.

I have a dataclass which already define `__getitem__`, `__iter__` for numpy-like operation, so I cannot support `collections.abc.Mapping`.
I would like to use it with tree too, so `chex.assert_xyz` works.