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.78k
stars
610
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
Added support of NANOO fp8.
#3993
wenchenvincent
closed
1 day ago
7
added `self.param` to `nnx.compat`
#3992
chiamp
opened
2 weeks ago
1
Added `nnx.graph.iter_children`
#3991
chiamp
closed
2 weeks ago
1
New python venv fails to pip install the flax mnist example requirements
#3990
sycamoreoak
closed
13 hours ago
3
[nnx] vectorize vmap split counts
#3989
cgarciae
closed
2 weeks ago
0
[nnx] LoRAParam inherits from Param
#3988
cgarciae
closed
1 day ago
0
updated `nnx.variables` docstrings
#3986
chiamp
closed
2 weeks ago
1
Adding Counter metric.
#3985
copybara-service[bot]
opened
3 weeks ago
0
[NVIDIA] Rename fp8 custom dtype to `fp32_max_grad`
#3984
kaixih
opened
3 weeks ago
3
added `wrt` option to `nnx.Optimizer`
#3983
chiamp
closed
2 weeks ago
1
Fixes for batch norm docs
#3982
jkarwowski
closed
3 weeks ago
1
fix deprecation warning
#3981
chiamp
closed
3 weeks ago
1
updated NNX `rnglib` docstring
#3980
chiamp
closed
3 weeks ago
1
Not sure if flax use the GPU
#3979
hndrbrm
closed
2 weeks ago
2
* Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported.
#3978
copybara-service[bot]
closed
3 weeks ago
0
Cannot load checkpoint saved in flax 0.5.3 with flax 0.6.1
#3977
tpgh24
closed
3 weeks ago
0
make flax work with upcoming JAX change to tree_map (being more careful about
#3976
copybara-service[bot]
closed
3 weeks ago
0
updated `nnx.training`
#3975
chiamp
closed
3 weeks ago
1
updated `nnx.Conv` and `nnx.ConvTranspose`
#3974
chiamp
closed
3 weeks ago
1
Using `vmap` on the loss function change results
#3973
gduflo
opened
3 weeks ago
0
updated NNX nn docstrings
#3972
chiamp
closed
3 weeks ago
1
updated nnx attention and nn docstrings
#3971
chiamp
closed
3 weeks ago
0
Modify Welford metric to return mean value.
#3970
copybara-service[bot]
closed
3 weeks ago
0
[nnx] add pmap
#3969
cgarciae
closed
3 weeks ago
1
Replace `XLACompatibleSharding` with `jax.sharding.Sharding` since the former is deprecated and will be removed in the future.
#3968
copybara-service[bot]
closed
3 weeks ago
0
Replace `XLACompatibleSharding` with `jax.sharding.Sharding` since the former is deprecated and will be removed in the future.
#3967
copybara-service[bot]
closed
3 weeks ago
0
updated `nnx.module` docstrings
#3966
chiamp
closed
3 weeks ago
1
[linen] allow checkpoint to cache
#3965
copybara-service[bot]
opened
4 weeks ago
0
[nnx] make State generic
#3964
cgarciae
closed
3 weeks ago
0
[nnx] experimental transforms
#3963
cgarciae
opened
4 weeks ago
1
feature_request: support for tabulate/summary in the NNX API
#3962
AshishKumar4
opened
4 weeks ago
1
promote nnx to root-level import
#3961
chiamp
opened
4 weeks ago
1
fix HEAD
#3960
chiamp
closed
4 weeks ago
1
Adding Welford metric.
#3959
copybara-service[bot]
closed
4 weeks ago
0
updated `nnx.graph` docstrings
#3958
chiamp
closed
3 weeks ago
1
Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.
#3957
copybara-service[bot]
closed
4 weeks ago
0
Inconsistent results when module is a property of another
#3956
epignatelli
closed
3 weeks ago
2
`optax.multi_transform` + `nnx.State`/`nnx.Optimizer` troubles
#3955
cgarciae
closed
3 weeks ago
0
Minor grammar fixes to NNX documentation.
#3953
mcsmart76
closed
4 weeks ago
1
Rename all the NNX tests to internal naming & build conventions.
#3952
copybara-service[bot]
closed
1 month ago
0
imagenet example needs a refreshed requirements.txt
#3950
sycamoreoak
opened
1 month ago
1
[nnx] refactor vmap
#3949
cgarciae
closed
3 weeks ago
1
Add direct penzai.treescope support for NNX objects.
#3948
copybara-service[bot]
closed
1 month ago
0
[nnx] improve vmap axis size detection
#3947
cgarciae
closed
1 month ago
0
static_argnums argument to flax.linen.remat not working as expected
#3946
dionhaefner
opened
1 month ago
2
nn.Dense doesn't respect the provided dtype
#3945
dfdx
closed
1 month ago
2
Add more activation functions from jax.nn.
#3944
carlosgmartin
opened
1 month ago
0
Update NNX installation instructions in mnist tutorial
#3943
Davidnet
closed
6 days ago
2
[nnx] fix nnx_basics dependencies
#3942
cgarciae
closed
1 month ago
1
v0.8.5
#3941
cgarciae
closed
1 month ago
0
Previous
Next