issues
search
google
/
flax
Flax is a neural network library for JAX that is designed for flexibility.
https://flax.readthedocs.io
Apache License 2.0
5.79k
stars
609
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
[nnx] improve filters guide
#4059
cgarciae
opened
11 hours ago
1
Is there anyway to analyze activations in flax?
#4058
mohamad-amin
opened
1 day ago
1
[linen] generalize transform caching
#4057
copybara-service[bot]
opened
1 day ago
0
Allow metadata pass-through in flax.struct.field
#4056
cool-RR
opened
1 day ago
0
Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file.
#4055
copybara-service[bot]
opened
2 days ago
0
Opaque XLA crash when initializing model
#4054
fernandopalafox
closed
2 days ago
1
Unsets `profile_duration_ms` in Flax ResNet example to prioritize `num_profile_steps`.
#4052
copybara-service[bot]
opened
3 days ago
0
Add functool.wraps() annotation to flax.nn.jit.
#4051
copybara-service[bot]
closed
3 days ago
0
Update nnx_basics.md
#4050
bmucsanyi
opened
3 days ago
2
[nnx] fix Variable overloads and add shape/dtype properties
#4049
cgarciae
closed
3 days ago
1
Fix typo in `nnx_basics` doc
#4047
rajasekharporeddy
closed
3 days ago
3
typo in nnx_basics.md
#4046
notnot
closed
3 days ago
3
Significant performance difference of NNX relative to equinox
#4045
jlperla
opened
5 days ago
15
allow imagenet training profiling to be disabled in config
#4043
copybara-service[bot]
closed
1 week ago
0
Port gemma sampling tutorial
#4041
copybara-service[bot]
opened
1 week ago
1
v0.8.6
#4040
cgarciae
closed
1 week ago
0
Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file.
#4039
copybara-service[bot]
closed
2 days ago
0
upgrade python to 3.10 + use pyupgrade
#4038
cgarciae
closed
1 week ago
0
Port gemma/sampler to NNX
#4037
copybara-service[bot]
closed
1 week ago
0
Update orbax_upgrade_guide.rst for async checkpointing usage examples
#4036
kaushaladiti-2802
opened
1 week ago
1
Feature request: Mixture of Experts example
#4034
SamKG
closed
1 week ago
1
lstm error
#4032
layssi
opened
1 week ago
1
[nnx] stabilize unsafe_pytree
#4030
cgarciae
opened
1 week ago
3
Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.
#4029
copybara-service[bot]
closed
1 week ago
0
[nnx] add Using Filters guide
#4028
cgarciae
closed
1 week ago
5
Suboptimal default initialization of q/k/v projections in `nn.MultiHeadDotProductAttention`
#4027
MasterSkepticista
opened
1 week ago
2
Call user-defined variable transforms before determining axis size in nn.vmap.
#4026
copybara-service[bot]
closed
1 week ago
0
Fix Gemma test to unbreak head
#4025
IvyZX
closed
1 week ago
1
Unpin the Orbax pip version.
#4024
copybara-service[bot]
closed
1 week ago
0
flax nn.tabulate Incorrectly Reports FLOPs and VJP FLOPs
#4023
Surya-77
opened
1 week ago
2
[nnx] add jit donation test
#4022
cgarciae
opened
1 week ago
0
[nnx] add Filters guide
#4021
cgarciae
closed
1 week ago
4
Port gemma/transformer to NNX
#4019
copybara-service[bot]
closed
1 week ago
0
Support Python 3.9
#4018
copybara-service[bot]
closed
2 weeks ago
0
Internal change to fix current head CI
#4017
copybara-service[bot]
closed
2 weeks ago
0
Unpin orbax version after their new release
#4016
IvyZX
closed
1 week ago
0
`save_checkpoint` fails with the most recent orbax release
#4015
apaleyes
closed
2 weeks ago
2
Port gemma/modules to NNX
#4014
copybara-service[bot]
closed
2 weeks ago
0
Port gemma/layers to NNX
#4013
copybara-service[bot]
closed
2 weeks ago
0
Helper function for loading params from a linen module
#4012
copybara-service[bot]
closed
2 weeks ago
0
[nnx] add PathContains Filter
#4011
cgarciae
closed
2 weeks ago
0
Will nnx.MultiHeadAttention support flash att?
#4009
Mddct
closed
2 weeks ago
3
Force `fp32` in `attention.MultiHeadDotProductAttention` for softmax operator
#4008
willisma
closed
1 week ago
1
[nnx] fix grad
#4007
cgarciae
closed
2 weeks ago
0
Fix head pytest breakage
#4006
IvyZX
closed
2 weeks ago
0
Add NNX surgery guide
#4005
IvyZX
closed
1 week ago
2
[nnx] pure API
#4004
cgarciae
opened
2 weeks ago
0
[nnx] remove display
#4003
cgarciae
opened
2 weeks ago
1
Fix pickling of exceptions
#4002
sanderland
closed
1 week ago
2
moved `Module.iter_*` methods to `nnx.graph`
#4001
chiamp
opened
2 weeks ago
0
Next