google/flax: Flax is a neural network library for JAX that is designed for flexibility.

3557
STARS
74
WATCHERS
408
FORKS
111
ISSUES

flax Recent Issues

Issue Title State Comments Created Date Updated Date
Unhelpful error message and a limitation of __post_init__ open 0 2022-09-23 2022-09-29
from_state_dict is too lenient open 2 2022-09-19 2022-09-29
Docs: Bug in Managing Parameters and State example closed 1 2022-09-14 2022-09-29
Add tests for pooling.py closed 1 2022-09-14 2022-09-21
[Accessibility] Enable EPUB output on ReadTheDocs open 0 2022-09-14 2022-09-29
Tensorflow Required Dependancy For Saving Checkpoint? open 1 2022-09-12 2022-09-29
Distributed training on Cloud TPUs stuck at 0% closed 0 2022-09-10 2022-09-29
flax.linen.remat with concrete=True doesn't work with jax 0.3.17 closed 1 2022-09-10 2022-09-29
Adding "count_include_pad" options to AvgPool layer open 4 2022-09-08 2022-09-29
Async save checkpoints not working closed 3 2022-09-07 2022-09-29
bug while fitting a multivariate normal closed 1 2022-09-07 2022-09-29
serialization.from_bytes not checking shapes of loaded arrays w.r.t. target open 1 2022-09-05 2022-09-29
FLIP: Generalized metadata open 0 2022-09-05 2022-09-29
Support serialization for `jax.tree_util.Partial` open 1 2022-09-05 2022-09-29
pmap got inconsistent sizes for array axes to be mapped: closed 1 2022-09-04 2022-09-29
Re-introduce pre-Linen JAXPR based lazy init into Linen open 0 2022-09-01 2022-09-29
Unnecessary assert in `merge_param` closed 2 2022-08-29 2022-09-29
Custom node type mismatch: expected type: <class 'flax.core.frozen_dict.FrozenDict'>, value: Traced<ShapedArray(float32[256,14,14,3])>with<DynamicJaxprTrace(level=0/1)>. closed 0 2022-08-29 2022-09-29
BatchNorm does not correctly compute mean/variance across workers when they have different batch sizes closed 4 2022-08-25 2022-09-21
`__getattr__` and pytype checking open 12 2022-08-24 2022-09-21
Test TransformedMethodReturnValueError for a returned Variable open 0 2022-08-23 2022-09-21
Can we use TrainState with pad_shard_upad? closed 2 2022-08-22 2022-09-21
Outdated `rich` dependency version closed 0 2022-08-18 2022-09-21
imagenet demo not working out of the box closed 5 2022-08-16 2022-09-21
Flax avg_pool also includes padding tokens when computing the mean open 1 2022-08-16 2022-09-21
Add a reference implementation for Wide ResNet closed 2 2022-08-16 2022-09-21
Half precision in flax models closed 1 2022-08-08 2022-09-21
Transpose Convolution module issue when used with circular padding and vmap closed 1 2022-08-02 2022-09-21
`nn.tabulate()` unexpectedly prints to the jupyter console closed 2 2022-08-02 2022-09-21
Some links on RTD do not work closed 1 2022-08-02 2022-09-21
ImportError: cannot import name 'isin' from closed 0 2022-08-01 2022-09-18
Scope.path names vs variables structure inconsistencies closed 1 2022-08-01 2022-09-18
Program suddenly hangs closed 0 2022-08-01 2022-09-18
Inconsistent type annotation of `step` in `training.checkpoints` closed 0 2022-07-27 2022-09-07
Pip install fails on macOS with Apple M1 and Python >= 3.9 closed 4 2022-07-27 2022-09-18
SORRY This page does not exist yet / 404 on google query and logs closed 1 2022-07-26 2022-09-18
Improve Landing Page example open 5 2022-07-25 2022-09-06
Improve Example section on RDT closed 0 2022-07-25 2022-09-17
System memory leak in Colab closed 4 2022-07-24 2022-09-12
Fix jax.tree_blahblah FutureWarning warnings due to jax symbol cleanup. closed 0 2022-07-23 2022-09-04
Typo in the PPO example closed 2 2022-07-19 2022-08-19
Contributing policy for Community Examples open 0 2022-07-18 2022-08-29
Simple LSTM example closed 0 2022-07-15 2022-09-23
Lost scope mutations after a transform on self closed 2 2022-07-14 2022-08-19
IPython autoreload is not supported by flax open 0 2022-07-11 2022-08-19
Would you like to add link to tutorial demonstrating Wav2Vec2 finetuning open 1 2022-07-10 2022-08-15
How to use nn.BatchNorm??? closed 0 2022-07-07 2022-08-15
Tuple[SomeType] type annotations should be Tuple[SomeType, ...] closed 0 2022-07-07 2022-08-15
`nn.tabulate` results in `KeyError: 'intermediates'` with methods that include transformations closed 2 2022-07-07 2022-09-04
linen.scan still have problem in latest release closed 1 2022-07-06 2022-08-08

google's Other Repos