Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add dtypes to the rest of eqx.nn #689

Merged
merged 2 commits into from
Mar 26, 2024

Conversation

Artur-Galstyan
Copy link
Contributor

Basically the same as #680 but for the rest of the API.

Double check SpectralNorm - I'm not 100% confident in that one.

Copy link
Owner

@patrick-kidger patrick-kidger left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On SpectralNorm, I think the appropriate thing to do is probably to use weight.dtype instead.

By the way, I think BatchNorm could also do with updating: it has a dtype argument but so far only uses these for the moving statistics, but not its learnt parameters. We should probably change that. Would you be able to do that as part of this PR?

Other than that, thank you for doing this! This is all great to see. :D

@Artur-Galstyan
Copy link
Contributor Author

Yup, done.

@patrick-kidger patrick-kidger merged commit eb4a577 into patrick-kidger:dev Mar 26, 2024
2 checks passed
@patrick-kidger
Copy link
Owner

Neat! LGTM, and thank you for the contribution. We'll have this appear in the next version of Equinox.

@patrick-kidger patrick-kidger mentioned this pull request Apr 14, 2024
patrick-kidger pushed a commit that referenced this pull request Apr 14, 2024
* added dtypes

* fixed norms and added info to _spectral_norm
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants