Penzai as a training framework
rosinality opened this issue · comments
Hello, thank you for a great framework!
I am currently considering penzai for building model trainer, as an alternative to frameworks like Flax. Would be there any possible limitations for penzai for this purpose? For example, using NamedArray would incur zero overheads? I think they are, but I want to be sure before proceeding with development.
Thank you!
Glad you like Penzai! As long as you run your code under jax.jit, I don't think there will be any overhead to using Penzai relative to frameworks like Flax.
(That said, Penzai's neural network standard library is pretty small right now, and focused on readability. So it may not have the same optimizations that some other frameworks support out of the box.)
If you do run into any issues with training models using Penzai, please let us know!