Revise 'HaikuModel' backend.

* Add proper weights and gradients naming.
* Store the model's nested dict of parameters rather than its
  flattened counterpart.
* Refactor some backend code to improve readability of gradients
  computation.
* Add a public getter to access weights' names, opt. restricted
  to trainable ones.
* Revise `set_trainable_weights` to received lists of names rather
  than indices. Keep support for predicates and PyTree inputs.