google / flax

Flax is a neural network library for JAX that is designed for flexibility.

google/flax Issues

`pmap` state
Updated 4 days ago
Add two aliases to activations import
Updated 4 days ago
Write a guide, HOWTO and/or example that shows how to correctly deal with the last batch during eval
Closed 5 days ago3
Switch to JAX's newer name-stack based named_call implementation by default.
Updated 5 days ago
HOWTOS tracking issue
Updated 5 days ago
Add HOWTO explaining how to load MNIST from torchvision and HuggingFace datasets
Closed 5 days ago1
Add (auto-generated) Colabs to all HOWTOs
Closed 5 days ago3
Automatically generate Colabs from all HOWTOs
Updated 5 days ago2
Add HOWTO: Data-parallel training
Closed 5 days ago3
Rewrite the Linen Upgrade Guide doc as a HOWTO
Closed 5 days ago1
Track Github metrics over time
Updated 10 days ago5
Improve Early Stopping API
Updated 11 days ago
We broke the 1:1 correspondence with attribute names and variable dict names
Updated 12 days ago3
Start using positional-only parameters in Linen
Updated 15 days ago1
`traverse_util.flatten_dict` function works differently in two different python environments despite same version of flax being installed
Closed 17 days ago2
Add HOWTO explaining how to set up fully reproducible experiments
Closed 19 days ago1
Impure functions used in documentation without clarification
Closed a month ago4
Import all activation functions from linen.activations in linen.__init__
Closed 19 days ago2
TensorFlow is not listed as a `setup.py` dependency but is required
Updated 20 days ago5
Annotated MNIST: The implementation of `cross_entropy_loss` is incorrect
Closed 20 days ago11
How to Implement MaxPooling2D with Ceil Mode in Flax?
Closed 20 days ago
Remove tree_multimap filterwarning once Optax releases a new version
Updated 20 days ago
Extracting Backbone Intermediates for Use in a FPN
Closed 20 days ago
Recommended Approach for Gradient Accumulation
Updated 20 days ago
Updating subtree with `put_variable` doesn't update sub-scopes' references.
Updated 20 days ago
Better document initialization functions
Updated 20 days ago
Document default stride for pooling functions
Closed 22 days ago
nn.scan + nn.BatchNorm
Closed 22 days ago3
integration with mup?
Closed 22 days ago
Adding a new transposed convolution function (similar to `torch.nn.ConvTranspose2d()`)
Updated a month ago8
UNKNOWN error when calling Conv
Closed a month ago3
What is `replace` in `TrainState`
Closed a month ago2
msgpack_serialize does not work with FrozenDict PyTree
Closed a month ago2
flax.linen.Conv needs better error checking of 'padding' argument.
Closed 2 months ago
How is `MultiHeadDotProductAttention` implemented?
Closed 2 months ago2
Can not assign dict whos key is not string as module attribute
Closed 2 months ago3
Clarify expectations for `variables` dict in apply()
Closed 2 months ago1
Create Colab Snippets
Updated 2 months ago
Improve documentation in flax.core
Updated 2 months ago2
Expand lifted custom VJP
Closed 2 months ago8
Improve lifted transforms docstrings
Updated 2 months ago1
Flax much slower than pure Jax?
Closed 2 months ago2
`typing.get_type_hints()` is broken for linen modules
Closed 3 months ago1
Incompatible variables for Tensorboard hparams are recast to strings but never returned
Closed 3 months ago1
How to extract intermediate values from "underneath" `vmap`?
Closed 3 months ago
Feature request: ability to apply stop gradient to some parameters
Closed 3 months ago7
Add lifted value_and_grad?
Closed 3 months ago
Crash when using generic module class
Closed 3 months ago
Opaque error message
Closed 3 months ago1
Add a HOWTO showing how to get per-example gradients
Closed 3 months ago1

ezoic increase your site revenue