Added colab notebook for state params#1932
Added colab notebook for state params#1932copybara-service[bot] merged 8 commits intogoogle:mainfrom
Conversation
|
Could you please take a look at this one? TIA |
|
Thanks a lot @AdityaKane2001! Would it be possible to also include the Colab itself in this PR, and add it to I also found that the Colab doesn't work, you need to change to dummy input to: dummy_input = jnp.ones((32, 5)) |
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
Could you please take a look and let me know it this is what you wanted? |
marcvanzee
left a comment
marcvanzee
left a comment
There was a problem hiding this comment.
Thanks! Just a few small comments.
docs/howtos/state_params.rst
Outdated
| @@ -1,3 +1,6 @@ | |||
| .. image:: https://colab.research.google.com/assets/colab-badge.svg | |||
| :target: https://colab.research.google.com/gist/AdityaKane2001/cbd09c2fcb86163cb301076cb60ce200/state_params_howto.ipynb | |||
There was a problem hiding this comment.
Please point to the Colab in this repo
| "opt_state = tx.init(params)\n", | ||
| "\n", | ||
| "for _ in range(num_epochs):\n", | ||
| " opt_state, params, state, loss = update_step(\n", |
There was a problem hiding this comment.
It would be nice to print something. maybe print the loss after each epoch?
There was a problem hiding this comment.
Sure, I didn't add that earlier because I thought it would be inconsistent with the docs.
|
It seems you cannot use |
|
Should I keep |
I think you can keep that, since it is in a comment. |
marcvanzee
left a comment
marcvanzee
left a comment
There was a problem hiding this comment.
Looks great, thanks a lot!
|
Thanks for the merge! |
2 participants
Regarding #1910
Added "Open in Colab" button which points to a Colab notebook having executable code of the Managing Parameters and State guide.
Checklist
checks if that's the case).
discussion (please add a
link).
documentation guidelines.
(No quality testing = no merge!)