flax
[nnx] always fold_in on fork + new ForkedKeys return type
#3722

Merged

[nnx] always fold_in on fork + new ForkedKeys return type #3722

copybara-service merged 1 commit into main from nnx-dynamic-rngs

cgarciae force pushed 1 year ago

cgarciae changed the title ~~[nnx] dynamic RngStream.count + ForkRngs~~ [nnx] always fold_in on fork + ForkRngs 1 year ago

cgarciae force pushed 1 year ago

cgarciae changed the title ~~[nnx] always fold_in on fork + ForkRngs~~ [nnx] always fold_in on fork + new ForkedKeys return type 1 year ago

levskaya approved these changes on 2024-02-28

cgarciae added pull ready

cgarciae force pushed 1 year ago

[nnx] dynamic RngStream.count + ForkRngs

3e693506

cgarciae force pushed to 3e693506 1 year ago

copybara-service merged acba0bf6 into main 1 year ago

copybara-service deleted the nnx-dynamic-rngs branch 1 year ago

Reviewers

levskaya

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone