Make MALA work with any PyTree

[rlouf]

Closes #214

[codecov]

Codecov Report

Merging #216 (d1ad9d7) into main (07c345a) will increase coverage by 0.00%.
The diff coverage is 100.00%.

❗ Current head d1ad9d7 differs from pull request most recent head 597657e. Consider uploading reports for the commit 597657e to get more accurate results

@@           Coverage Diff           @@
##             main     #216   +/-   ##
=======================================
  Coverage   98.32%   98.33%           
=======================================
  Files          40       40           
  Lines        1616     1618    +2     
=======================================
+ Hits         1589     1591    +2     
  Misses         27       27           

Impacted Files	Coverage Δ
blackjax/mcmc/mala.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 07c345a...597657e. Read the comment docs.

[junpenglao]

Add a small test?