google flax issues - Githubissues

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 608 forks source link

issues

Newest

Newest Most commented Recently updated Oldest Least commented Least recently updated

Significant performance difference of NNX relative to equinox

#4045 jlperla opened 23 hours ago
2
allow imagenet training profiling to be disabled in config

#4043 copybara-service[bot] closed 2 days ago
0
Port gemma sampling tutorial

#4041 copybara-service[bot] opened 3 days ago
1
v0.8.6

#4040 cgarciae closed 2 days 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] opened 3 days ago
0
upgrade python to 3.10 + use pyupgrade

#4038 cgarciae closed 3 days ago
0
Port gemma/sampler to NNX

#4037 copybara-service[bot] closed 4 days ago
0
Update orbax_upgrade_guide.rst for async checkpointing usage examples

#4036 kaushaladiti-2802 opened 4 days ago
1
Feature request: Mixture of Experts example

#4034 SamKG closed 3 days ago
1
lstm error

#4032 layssi opened 4 days ago
0
[nnx] stabilize unsafe_pytree

#4030 cgarciae opened 5 days ago
3
Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.

#4029 copybara-service[bot] closed 5 days ago
0
[nnx] add Using Filters guide

#4028 cgarciae closed 2 days ago
5
Suboptimal default initialization of q/k/v projections in `nn.MultiHeadDotProductAttention`

#4027 MasterSkepticista opened 6 days ago
2
Call user-defined variable transforms before determining axis size in nn.vmap.

#4026 copybara-service[bot] closed 5 days ago
0
Fix Gemma test to unbreak head

#4025 IvyZX closed 6 days ago
1
Unpin the Orbax pip version.

#4024 copybara-service[bot] closed 6 days ago
0
flax nn.tabulate Incorrectly Reports FLOPs and VJP FLOPs

#4023 Surya-77 opened 6 days ago
2
[nnx] add jit donation test

#4022 cgarciae opened 1 week ago
0
[nnx] add Filters guide

#4021 cgarciae closed 6 days ago
4
Port gemma/transformer to NNX

#4019 copybara-service[bot] closed 4 days ago
0
Support Python 3.9

#4018 copybara-service[bot] closed 1 week ago
0
Internal change to fix current head CI

#4017 copybara-service[bot] closed 1 week ago
0
Unpin orbax version after their new release

#4016 IvyZX closed 6 days ago
0
`save_checkpoint` fails with the most recent orbax release

#4015 apaleyes closed 1 week ago
2
Port gemma/modules to NNX

#4014 copybara-service[bot] closed 1 week ago
0
Port gemma/layers to NNX

#4013 copybara-service[bot] closed 1 week ago
0
Helper function for loading params from a linen module

#4012 copybara-service[bot] closed 1 week ago
0
[nnx] add PathContains Filter

#4011 cgarciae closed 1 week ago
0
Will nnx.MultiHeadAttention support flash att?

#4009 Mddct closed 1 week ago
3
Force `fp32` in `attention.MultiHeadDotProductAttention` for softmax operator

#4008 willisma closed 5 days ago
1
[nnx] fix grad

#4007 cgarciae closed 1 week ago
0
Fix head pytest breakage

#4006 IvyZX closed 1 week ago
0
Add NNX surgery guide

#4005 IvyZX closed 4 days ago
2
[nnx] pure API

#4004 cgarciae opened 1 week ago
0
[nnx] remove display

#4003 cgarciae opened 1 week ago
1
Fix pickling of exceptions

#4002 sanderland closed 6 days ago
2
moved `Module.iter_*` methods to `nnx.graph`

#4001 chiamp opened 1 week ago
0
Exceptions are not pickle-able

#4000 sanderland closed 6 days ago
0
`serialization.from_state_dict` does not restore to jax.Arrays

#3999 PhilipVinc opened 2 weeks ago
1
NNX + mutable state + JIT = "Cannot mutate <module name> from a different trace level"

#3997 dfdx closed 2 weeks ago
0
CI: add test run against oldest supported jax version

#3996 jakevdp closed 5 days ago
1
[nnx] fix vmap

#3995 copybara-service[bot] closed 2 weeks ago
0
Best practice of dealing with sporadic FrozenDict conversions?

#3994 TimSchneider42 opened 2 weeks ago
2
Added support of NANOO fp8.

#3993 wenchenvincent opened 2 weeks 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 opened 2 weeks ago
3
[nnx] vectorize vmap split counts

#3989 cgarciae closed 2 weeks ago
0
[nnx] LoRAParam inherits from Param

#3988 cgarciae opened 2 weeks ago
0