-
Notifications
You must be signed in to change notification settings - Fork 635
Issues: google/flax
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
Author
Label
Projects
Milestones
Assignee
Sort
Issues list
nnx.jit(aux_fn) is slower than directly using nnx.jit(model.__call__)
#4218
opened Sep 23, 2024 by
JunhongXu
Depthwise convolution in
flax.nnx.Conv
is significantly slower than PyTorch and TensorFlow
#4207
opened Sep 18, 2024 by
YushaArif99
nnx.vmap
use the same random key rngs
inside nnx.Module across vectorization.
#4195
opened Sep 15, 2024 by
maxencefaldor
MultiHeadAttention documentation missing descriptions of
None
values for optional arguments
#4182
opened Sep 9, 2024 by
carlosgmartin
Support for optax lbfgs and related optimizers with NNX
Priority: P1 - soon
Response within 5 business days. Resolution within 30 days. (Assignee required)
#4144
opened Aug 26, 2024 by
jlperla
nn.WeightNorm needs a special scale_init to match PyTorch weight_norm
#4138
opened Aug 21, 2024 by
DBraun
Flax dataclasses do not support new-style PRNGKeyArray in serialisation
Priority: P1 - soon
Response within 5 business days. Resolution within 30 days. (Assignee required)
#4130
opened Aug 19, 2024 by
PhilipVinc
SPMD for initializing model using nnx.jit
Priority: P1 - soon
Response within 5 business days. Resolution within 30 days. (Assignee required)
#4129
opened Aug 18, 2024 by
mmorinag127
Clarification for LSTMCell Documentation
Priority: P2 - no schedule
Best effort response and resolution. We have no plan to work on this at the moment.
#4124
opened Aug 12, 2024 by
corentinlger
DynamicScale
behaves unexpected when computing per-sample gradients with vmap
.
#4114
opened Aug 2, 2024 by
hlzl
Previous Next
ProTip!
Add no:assignee to see everything that’s not assigned.