Optimizer should support pytree-based parameters #54
Description
The Optimizer, ParamSpace, and related components should operate on pytrees instead of flat dictionaries. This would enable better JAX compatibility, support nested parameter structures, simplify tree operations (fit, init, sample), and improve flexibility for complex domains. Padding logic, transformation handling, and OptimizerState storage should be updated accordingly.