flax
[nnx] RNN: add broadcast_rngs and state_axes APIs
#4407

Merged

[nnx] RNN: add broadcast_rngs and state_axes APIs #4407

copybara-service merged 1 commit into main from rnn-broadcast-rngs

cgarciae force pushed 1 year ago

[nnx] RNN: add broadcast_rngs and state_axes APIs

b9f016a0

cgarciae force pushed to b9f016a0 1 year ago

IvyZX approved these changes on 2025-01-08

IvyZX added pull ready

copybara-service merged 5408c40c into main 354 days ago

copybara-service deleted the rnn-broadcast-rngs branch 354 days ago

Reviewers

IvyZX

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone