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
6.15k
stars
648
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
Slow training step occasionally due to slow graph flatten
#4336
kriscao-cohere
opened
4 weeks ago
3
Remove remaining references to Penzai from NNX.
#4335
copybara-service[bot]
closed
4 weeks ago
1
[nnx] cleanup gemma notebook
#4334
cgarciae
closed
3 weeks ago
2
[nnx] improve transforms guide
#4333
cgarciae
closed
3 weeks ago
1
Avoid assert_array_equal on PRNG keys.
#4332
jakevdp
closed
4 weeks ago
0
fix ruff complaints
#4331
levskaya
closed
4 weeks ago
0
nnx.pmap documentation seems to be for vmap
#4330
hrbigelow
opened
4 weeks ago
2
nnx.jit recompiles for new instance of Module even if parameter shapes and dtypes are unchanged
#4329
kriscao-cohere
closed
4 weeks ago
3
get wrong structure when using jax2tf to convert nnx.module into tflite file
#4328
noahzhy
closed
4 weeks ago
1
Add support to jax2tf when using nnx.Module
#4327
noahzhy
closed
4 weeks ago
1
Avoid depending on JAX internals, which are about to change.
#4326
copybara-service[bot]
closed
4 weeks ago
0
Upgrade Flax NNX Gemma Sampling Inference doc
#4325
8bitmp3
opened
1 month ago
5
Remove outdated compatibility code.
#4324
jakevdp
closed
4 weeks ago
0
Delete the `nnx.bridge.module` code because we no longer go down that third module class path.
#4323
copybara-service[bot]
closed
1 month ago
0
remove markdown from section titles
#4322
cgarciae
closed
4 weeks ago
1
Avoid depending on JAX internals, which are about to change (see cl/681582933).
#4321
copybara-service[bot]
opened
1 month ago
0
Extracting learning rate from nnx.optimizer directly
#4320
noahzhy
closed
3 weeks ago
2
Is any method to call tabulate_fn on nnx.Module?
#4319
noahzhy
closed
1 month ago
1
[nnx] try casting integers keys in State.replace_by_pure_dict
#4317
cgarciae
closed
1 month ago
0
[nnx] improve mnist tutorial
#4316
cgarciae
closed
1 month ago
1
Avoid using float32 in normalization for mean/var and scale/bias parameters when force_float32_reductions=False
#4314
copybara-service[bot]
closed
4 weeks ago
0
Fix philosophy link
#4313
jorisSchaller
closed
1 month ago
2
Bug - Cannot access attribute "replace" for class "c123"
#4312
MRiabov
opened
1 month ago
5
Fixing numbering in mnist_tutorial.md
#4311
v1vekkumar
opened
1 month ago
4
Can't save checkpoint with orbax when using zero-size parameters
#4309
dionhaefner
opened
1 month ago
3
Issue on restoring NNX state from Orbax checkpoint as pure dict
#4308
mmorinag127
closed
1 month ago
2
Avoid float32 promotion in BatchNorm when using_running_average=True and force_float32_reductions=False
#4307
copybara-service[bot]
closed
1 month ago
0
[nnx] fix custom_vjp
#4306
cgarciae
closed
3 weeks ago
2
Flesh out the Haiku/Flax guide
#4305
IvyZX
closed
1 month ago
1
Remove GeGLU activation function and golden tests.
#4303
copybara-service[bot]
closed
4 weeks ago
0
Add Flax NNX GraphDef docstring
#4302
8bitmp3
closed
1 month ago
1
How to use nnx.custom_vjp with non-class arguments? Example needed
#4301
hrbigelow
closed
1 week ago
3
Fix devcontainer setup
#4299
jorisSchaller
closed
1 month ago
1
Dev container is outdated
#4298
jorisSchaller
closed
1 month ago
0
Update Flax NNX Scale Up guide
#4296
8bitmp3
closed
1 month ago
1
NNX Conv Layer Input Tuple Error
#4295
riverliway
opened
1 month ago
1
Ugrade Flax NNX Checkpointing guide
#4294
8bitmp3
closed
1 month ago
1
[nnx] flaxlib flatten
#4293
cgarciae
opened
1 month ago
0
update version to 0.10.0
#4292
cgarciae
closed
1 month ago
0
[nnx] improve init_cache docs
#4291
cgarciae
closed
1 month ago
0
Update Flax Evolution from Linen to NNX guide
#4289
8bitmp3
closed
1 month ago
0
[nnx] state filters
#4288
cgarciae
closed
1 month ago
0
Update Flax NNX vs JAX Transformations guide
#4286
8bitmp3
opened
1 month ago
0
Remove notebook exceptions in `docs_nnx` doctest
#4285
IvyZX
closed
1 month ago
1
Update Flax NNX Glossary
#4284
8bitmp3
closed
1 month ago
1
Fix Philosophy links, update copyright year,
#4283
8bitmp3
opened
1 month ago
2
Fixed nnx examples bad links in the README.md
#4282
vfdev-5
closed
1 month ago
0
[nnx] improve randomness guide
#4281
cgarciae
closed
1 month ago
1
Update Flax NNX Randomness
#4279
8bitmp3
opened
1 month ago
4
[nnx] add PrefixMapping
#4278
cgarciae
closed
1 month ago
0
Previous
Next