google/flax

Issues

• SPMD for initializing model using nnx.jit

  #4129 opened 4 months ago by mmorinag127
  2

• Pickling a trained model (NNX)

  #4247 opened 3 months ago by eugene
  7

• ToLinen is not hashable (Linen modules are)

  #4156 opened 4 months ago by PhilipVinc
  3

• How to use max_pool with xla_gpu_deterministic_ops?

  #4233 opened 3 months ago by blackblitz
  0

• Add NNX support for legacy `jax.random.PRNGKey()`

  #4231 opened 3 months ago by cisprague
  1

• nnx.jit(aux_fn) is slower than directly using nnx.jit(model.__call__)

  #4218 opened 3 months ago by JunhongXu
  3

• How to freeze parameters with `nnx` and `optax`?

  #4167 opened 4 months ago by maxencefaldor
  7

• nn.scan messes with jax's sharding

  #4217 opened 3 months ago by sokol11
  0

• nnx.Swish, jax.swish,... change the input shape

  #4214 opened 3 months ago by leson207
  2

• `test_vmap_and_cond_passthrough[_error]` tests fail

  #4211 opened 3 months ago by GaetanLepage
  3

• Depthwise convolution in `flax.nnx.Conv` is significantly slower than PyTorch and TensorFlow

  #4207 opened 3 months ago by YushaArif99
  0

• `nnx.vmap` use the same random key `rngs` inside nnx.Module across vectorization.

  #4195 opened 3 months ago by maxencefaldor
  1

• Significant performance difference of NNX relative to equinox

  #4045 opened 6 months ago by jlperla
  17

• `TODO: add notes link` in flax_basics guide

  #4193 opened 3 months ago by garymm
  0

• flax counterpart for `torch.nn.Conv1d`

  #4188 opened 4 months ago by Liyang90
  0

• MultiHeadAttention documentation missing descriptions of `None` values for optional arguments

  #4182 opened 4 months ago by carlosgmartin
  0

• Contributing to Flax Community Tutorials

  #4181 opened 4 months ago by yilunyu
  0

• CHANGELOG has not been updated since 0.8.2

  #4179 opened 4 months ago by enolan
  0

• [FLIP] JAX-style NNX Transforms

  #4107 opened 5 months ago by cgarciae
  0

• bias and kernel params are put on different gpu devices

  #4116 opened 5 months ago by YunxiTang
  1

• Flax 0.9.0 broke nnx rng splitting

  #4153 opened 4 months ago by kasper0406
  2

• `nnx.Dropout` in combination with `jax.lax.scan` returns an error

  #4151 opened 4 months ago by maxencefaldor
  0

• Discrepency between `jax.jit` and `nnx.jit` for `static_argnames`

  #4146 opened 4 months ago by maxencefaldor
  2

• Experimental-pytree flag causes crash

  #4142 opened 4 months ago by NeilGirdhar
  2

• Support for optax lbfgs and related optimizers with NNX

  #4144 opened 4 months ago by jlperla
  1

• nn.WeightNorm needs a special scale_init to match PyTorch weight_norm

  #4138 opened 4 months ago by DBraun
  1

• Flax dataclasses do not support new-style PRNGKeyArray in serialisation

  #4130 opened 4 months ago by PhilipVinc
  0

• Add IndRNN

  #4133 opened 4 months ago by carlosgmartin
  1

• Truncated Normal initializer doesn't match PyTorch

  #4091 opened 5 months ago by DBraun
  2

• nnx.initializers.uniform should support custom lower and upper bounds

  #4127 opened 4 months ago by SamPruden
  3

• Clarification for LSTMCell Documentation

  #4124 opened 4 months ago by corentinlger
  2

• MNIST tutorial broken for Colab TPU

  #4122 opened 5 months ago by rcrowe-google
  0

• [nnx] How to access the nnx.Param which is inside of the Module

  #4117 opened 5 months ago by mmorinag127
  2

• Flax.linen.conv unexpected behavior.

  #4113 opened 5 months ago by NITHISHM2410
  1

• Edit parameters of flax module in another module

  #4109 opened 5 months ago by minhkhoi1026
  2

• `DynamicScale` behaves unexpected when computing per-sample gradients with `vmap`.

  #4114 opened 5 months ago by hlzl
  0

• Performance penalty of typecast in `nn.Embed.__call__`

  #4100 opened 5 months ago by MasterSkepticista
  0

• lowering / cost analysis of @nnx.jit functions

  #4094 opened 5 months ago by cgarciae
  1

• NNXWrapper

  #4088 opened 5 months ago by PhilipVinc
  2

• Dropout seems not compatible with jax.jit

  #4085 opened 5 months ago by richardmkit
  1

• Large Difference in Loss between JAX and FLAX Two-Layer Linear Autoencoder

  #4087 opened 5 months ago by yCobanoglu
  1

• GroupNorm missing from NNX normalization layer

  #4086 opened 5 months ago by treigerm
  2

• lstm error

  #4032 opened 6 months ago by layssi
  2

• Docs: please clarify how to vmap nnx.Module over batch dimension

  #4061 opened 6 months ago by bionicles
  3

• Is there anyway to analyze activations in flax?

  #4058 opened 6 months ago by mohamad-amin
  1

• Opaque XLA crash when initializing model

  #4054 opened 6 months ago by fernandopalafox
  1

• typo in nnx_basics.md

  #4046 opened 6 months ago by notnot
  3

• Feature request: Mixture of Experts example

  #4034 opened 6 months ago by SamKG
  1

• Suboptimal default initialization of q/k/v projections in `nn.MultiHeadDotProductAttention`

  #4027 opened 6 months ago by MasterSkepticista
  2

• flax nn.tabulate Incorrectly Reports FLOPs and VJP FLOPs

  #4023 opened 6 months ago by Surya-77
  2