Fabian Pedregosa

duplicated https://github.com/google-deepmind/optax/pull/1000 since I'm having some issues getting this merged due to internal errors

yup. Finally merged 🎉

@carlosgmartin : we could also use some help if you have cycles to look into this ...

hey @SauravMaheshkar , yes, I believe a short example on this on the example/ directory would be welcomed!. I'll assign you this issue. Thanks for suggesting this!

instead of a hard-coded 1e-12, could we perhaps replace it with `np.finfo(embeddings.dtype).eps` ? this way the epsilon will depend on the dtype of the embeddings (which I believe is what...

there were some internal errors and then i forgot about it. Taking a look into it now

thanks @amosyou ! Yes, that's fine, do it however way its most convenient for you 😄

@carlosgmartin , there are some conflicts with main, do you mind updating the pull request? thanks

hey @carlosgmartin , thanks for the ping and apologies for the late reply, most of the team is on vacation 🏖️ The addition of the axis argument looks good to...

excellent, thanks for the info!