- [ ] Toy model conversion from `PyTorch` - [ ] Toy model conversion from `keras` (i.e. show lazy approach) - [x] orbax for checkpointing