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
MNIST NNX Tutorial broken links
#4394
stergiosba
opened
11 hours ago
0
Automated Code Change
#4393
copybara-service[bot]
opened
1 day ago
0
Issue with Optimizer Update in A2C Network with Optax Body:
#4391
Tomato-toast
opened
2 days ago
4
Fix fori_loop and while_loop on multiple modules
#4390
IvyZX
closed
1 day ago
0
ResNet-18 Initialization
#4388
ZedongPeng
opened
2 days ago
0
update version to 0.10.2
#4387
cgarciae
closed
2 days ago
0
Upgrade Flax readme to NNX
#4386
8bitmp3
closed
1 day ago
0
Update `async_checkpointer.py` reference
#4385
emmanuel-ferdman
opened
6 days ago
1
[nnx] add performance guide notebook
#4384
cgarciae
opened
6 days ago
1
Flax NNX and Orbax Checkpointing require hacking to work together
#4383
hdrwilkinson
opened
1 week ago
2
JAX for the impatient was deleted, but it is still referenced twice
#4382
ArthurConmy
opened
1 week ago
1
[nnx] add checkify
#4381
cgarciae
closed
1 week ago
0
Setup the flaxlib in C++, using Meson and Nanobind.
#4380
copybara-service[bot]
closed
1 week ago
0
[nnx] fix while_loop/fori_loop bug when sharing references
#4379
cgarciae
closed
1 week ago
0
[nnx] add some optimizations to graph.py
#4377
cgarciae
closed
6 days ago
0
[nnx] use HashableMapping instead of FrozenDict
#4376
cgarciae
closed
1 week ago
0
7x7 `nnx.Conv` using `float32` parameter dtype overflows(?) to `nan` when sharded
#4375
joshhansen
opened
1 week ago
0
Add `flax.nnx.eval_shape` docstring
#4374
8bitmp3
closed
1 week ago
0
Add `flax.nnx.remat` docstring
#4373
8bitmp3
closed
1 week ago
1
Add `flax.nnx.value_and_grad` docstring
#4372
8bitmp3
opened
1 week ago
5
Lint flax.nnx.while_loop docstring
#4371
8bitmp3
closed
6 days ago
0
Lint flax.nnx.fori_loop docstring
#4370
8bitmp3
closed
6 days ago
0
Update flax.nnx.jit docstring
#4369
8bitmp3
opened
1 week ago
0
Update `flax.nnx.grad` docstring
#4368
8bitmp3
opened
1 week ago
0
Initialization of `flax.linen` model parameters depends on member variable names
#4367
bluesunb
opened
1 week ago
1
flax does not work with sparse matrices
#4366
Prashant-Shekhar-Rao
closed
1 week ago
1
Fix multiple links to Orbax documentation
#4364
Matt-Hurd
opened
1 week ago
1
Avoid assert_array_equal for PRNG keys.
#4363
copybara-service[bot]
closed
1 week ago
0
Update nnx_basics.ipynb
#4362
mohsenh17
closed
1 week ago
3
Remove this line once diffbase is submitted.
#4361
copybara-service[bot]
closed
2 weeks ago
0
Fix PRNG handling in `nn.jit` under `nn.scan`.
#4359
copybara-service[bot]
closed
2 weeks ago
1
Scaling issue on CPU (X86 and arm) while using models from huffing face.
#4357
choudhary-devang
opened
2 weeks ago
0
[nnx] add unsafe_graph_cache decorator
#4356
cgarciae
opened
2 weeks ago
0
nnx.vmap example from documentation raise an index error
#4355
jhn-nt
opened
2 weeks ago
5
[nnx] add data parallel toy example
#4354
cgarciae
closed
1 week ago
0
Add `nnx.fori_loop`
#4353
IvyZX
closed
3 weeks ago
0
[nnx] support pure dicts
#4352
cgarciae
closed
1 week ago
0
Linesearch (and lbfgs) support
#4351
jlperla
closed
2 weeks ago
5
Add logical axis global context support for NNX
#4350
IvyZX
closed
1 week ago
0
Define model surgery in docs
#4349
8bitmp3
closed
3 weeks ago
1
Add support for storing arbitrary PyTrees with `Module.perturb()`
#4348
copybara-service[bot]
closed
3 weeks ago
0
[nnx] Optimizer uses variables
#4347
cgarciae
opened
3 weeks ago
1
[nnx] update State and variables docstrings
#4346
cgarciae
closed
3 weeks ago
0
update version to v0.10.1
#4345
cgarciae
closed
3 weeks ago
0
Can @nnx.custom_vjp bwd function update its arguments?
#4344
hrbigelow
closed
3 weeks ago
2
Add NNX transforms `nnx.while_loop` and `nnx.switch`
#4343
IvyZX
closed
3 weeks ago
1
Docstring for State is wrong about leaf types (or else the FSDP example is wrong)
#4342
gabbard
closed
3 weeks ago
1
Remove non-lazy RNG compat mode and flag from flax.
#4339
copybara-service[bot]
closed
3 weeks ago
0
Add API reference for flax.nnx.nn and improve landing page
#4338
IvyZX
closed
3 weeks ago
0
Fix typos in Flax NNX Migrating from Haiku to Flax
#4337
8bitmp3
closed
3 weeks ago
1
Next