Feature Request: Penzai-like Library for PyKAN

[vrtnis]

Hi!

While working with PyKAN, I came across a library called Penzai which offers a highly legible and functional approach to working with JAX models.

Currently, I am working with Penzai in marimo which is an excellent reactive notebook platform.

I believe having similar features in PyKAN would significantly enhance its utility and user experience. Below are some key features from Penzai that could be highly beneficial if integrated into PyKAN:

Some benefits I can think of

• Enhanced Model Interpretability: Tools like an advanced pretty-printer and pytree selectors would make it easier to visualize and understand KAN models.
• Improved Debugging and Modifying: With declarative combinator-based libraries and advanced selector tools, users can easily debug and modify their models.
• Greater Flexibility: A named axis system would allow for more flexible and readable code, making it easier to manage complex models.

Thank you for pykan! Let me know if I can provide further context and examples!