Saurav Maheshkar

Even `mypy` and Google Style Guide can be checked IMO. Gentle ping @levskaya

Updated Tracker - [x] Port to linen API - [x] Update README.md - [x] Add requirements.txt - [x] Update file structure - [x] Use `ml_collections.ConfigDict` - [ ] Add benchmark...

This got lost in my inbox 😅 , if anyone else wants to take this up, feel free

@hwchase17 Request for Review

I can confirm that on Colab running the same commands as mentioned in the original issue comment imports without any error. ```python import jax !pip install --quiet flax import flax...

> Hi @SauravMaheshkar Thank you very much for your contribution! Looking forward to contributing more 😄

Looks like the tests passed 😄

> @marcvanzee To the best of my knowledge, the use of `flax.traverese_utils.Traversal` is now deprecated for `flax.optim`. So is there still the need of updating the `flax.traverese_utils` API? If yes,...

The `FlaxTrainStep` does have a `lr_scheduler` argument. I can't understand the issue can you provide more context ?? https://github.com/allenai/tango/blob/1ee2c56a0f557de3b2ba0bb25432487f0c4c50b5/tango/integrations/flax/train.py#L84

Failing tests seem unrelated to the PR