make notebook (colab) examples work on TPUs

[fabianp]

these are the examples on docs/notebooks/

Probably it's necessary to add this line when running on TPUs https://github.com/google/jax#pip-installation-colab-tpu

I think this would be a good first issue

[Samad999777]

Hey @fabianp is this still open for contribution ?

[fabianp]

yes!

[Samad999777]

Great, I would like to work on it

[fabianp]

all yours :-)

[Samad999777]

I actually have to reproduce the already given examples in colab notebooks using TPU?

[fabianp]

yes, the goal is that the examples run and produce similar results using TPU instead of GPU. Some might work out of the box, some might fail (nobody knows because nobody has tried yet)

[Samad999777]

Alright, I will keep you posted about it. You can assign this issue to me.

[Samad999777]

Hey @fabianp I have a colab notebook working for one example. Should i create a pull request so you can check it out?

[fabianp]

Yes please :-)

On Sat., May 28, 2022, 14:52 Samad999777, @.***> wrote:

Hey @fabianp https://github.com/fabianp I have a colab notebook working for one example. Should i create a pull request so you can check it out?

— Reply to this email directly, view it on GitHub https://github.com/google/jaxopt/issues/202#issuecomment-1140313428, or unsubscribe https://github.com/notifications/unsubscribe-auth/AACDZB7G2WSMCZDD3LOUC3LVMJTM5ANCNFSM5TK2DLQA . You are receiving this because you were mentioned.Message ID: @.***>

[fabianp]

@Samad999777 : any updates on this :-) ?

[Samad999777]

@fabianp im sorry i got positive for covid, im mostly recovered now. Give me 2 days to setup my laptop again and i'll push the changes.

[fabianp]

There's no rush, I was reaching out mostly in case you needed some help. Again, there's no urgency in this

On Mon., Jun. 20, 2022, 16:03 Samad999777, @.***> wrote:

@fabianp https://github.com/fabianp im sorry i got positive for covid, im mostly recovered now. Give me 2 days to setup my laptop again and i'll push the changes.

— Reply to this email directly, view it on GitHub https://github.com/google/jaxopt/issues/202#issuecomment-1160804723, or unsubscribe https://github.com/notifications/unsubscribe-auth/AACDZB7OFCNSEQOTFS7VAYDVQDFCJANCNFSM5TK2DLQA . You are receiving this because you were mentioned.Message ID: @.***>

[Samad999777]

@fabianp thanks alot, hopefully we will wrap it up this week.