Add a hamiltonian rule for the jax operators (#1514)

This is what is left of https://github.com/netket/netket/pull/968

 rebased on master.
---------

Co-authored-by: Filippo Vicentini <filippovicentini@gmail.com>