Skip to content

Document direct NFNet use with hk.transform - #753

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:masterfrom
sylvesterkaczmarek:docs/nfnet-haiku-transform
Open

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:masterfrom
sylvesterkaczmarek:docs/nfnet-haiku-transform

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown

Summary

Add a minimal example showing how to instantiate and run NFNet correctly through Haiku.

NFNet subclasses hk.Module, so constructing it directly at module scope raises ValueError: All hk.Modules must be initialized inside an hk.transform. The README previously pointed users to the experiment code and Colab but did not show the minimal pattern for direct model use.

This change:

  • constructs NFNet inside a function passed to hk.transform
  • shows parameter initialization and inference
  • notes that training-time application needs a PRNG key for stochastic layers such as dropout

Fixes #171.

Validation

Documentation-only change. The example matches the current NFNet.__init__ and NFNet.__call__ signatures and the Haiku transform contract.

The final branch is based on current upstream master, contains one commit, and changes only nfnets/README.md.

@polarbe

polarbe commented Aug 23, 2026 via email

Copy link
Copy Markdown

@sylvesterkaczmarek

Copy link
Copy Markdown
Author

Could a maintainer review the direct NFNet use with hk.transform documentation and arrange a fresh approved workflow run? The previous security run ended with no jobs after 30 days, which looks like expired approval.

@polarbe

polarbe commented Oct 7, 2026 via email

Copy link
Copy Markdown

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.

ValueError: All hk.Modules must be initialized inside an hk.transform.

2 participants