I’m not sure if I’ll be able to devote enough time but I will try. I am also planning to work on another project.
However, @ttj is planning to contribute Flax FNet to the repository and I am planning to work on PyTorch parallely. Wdyt?  Is it okay if we only work on Flax? or use the repo directly?