flax
1ccd1598 - Allow using LazyRngs for flax init/apply.

Commit

1 year ago

Allow using LazyRngs for flax init/apply. This is occasionally useful for refactoring while preserving rng derivations. PiperOrigin-RevId: 783598096

References

Author

levskaya

Committer

a-googler

Parents