[Enhancement] Investigate/use Chex for assertions and type handling

[atiyo]

The use of checkify.check can alter the return type of jitted functions, so it has been proposed to investigate chex (https://github.com/google-deepmind/chex) to replace the checkify.check calls in https://github.com/pasqal-io/horqrux/compare/feature/psr_on_all_gates

Also, it might be more streamlined in general to replace custom pytrees, e.g. in Primitive and Parametric among others with chex.dataclasses.