Optimizer should support pytree-based parameters

[alonfnt]

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.