Pull Requests google/flax

[nnx] add optax optimizer

#4752 opened 2025-05-16 04:50 by cgarciae

Remove VariableState

#4729 opened 2025-04-25 16:49 by cgarciae

FEAT: add LoRALinearGeneral

#4718 opened 2025-04-19 09:17 by zhengzl18

convert DenseGeneral to NNX

#4710 opened 2025-04-17 18:31 by copybara-service[bot]

Add `attention_bias` parameter to `MultiHeadDotProductAttention`.

#4701 opened 2025-04-13 09:28 by copybara-service[bot]

Added LSTM Model in JAX from scratch

#4690 opened 2025-04-05 21:49 by Tushaam

bump webhook action version and improve permissions

#4685 opened 2025-04-04 18:11 by ZacCranko

Internal change

#4682 opened 2025-04-04 01:27 by copybara-service[bot]

Bump pyo3 from 0.21.2 to 0.24.1 in /flaxlib_src in the cargo group across 1 directory dependencies rust

#4675 opened 2025-04-02 13:23 by dependabot[bot]

[nnx] refactor remat pull ready

#4662 opened 2025-03-28 00:38 by cgarciae

add experimental support for mutable array

#4661 opened 2025-03-28 00:36 by cgarciae

add fast_flatten

#4654 opened 2025-03-26 00:22 by cgarciae

roll back #4646

#4647 opened 2025-03-24 20:31 by copybara-service[bot]

Apply vmap jnp.pad over the batch in REFLECT convolution mode

#4640 opened 2025-03-20 09:49 by sarlinpe

[bridge module] Add `bridge.share_scope` for layer-sublayer pairs.

#4638 opened 2025-03-19 00:02 by copybara-service[bot]

Add DynamicTanh (https://arxiv.org/pdf/2503.10622)

#4631 opened 2025-03-17 22:23 by seanemoon

[nnx] add Pytree

#4625 opened 2025-03-15 00:06 by cgarciae

Add pytrees as graph nodes again.

#4605 opened 2025-03-07 04:43 by copybara-service[bot]

Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec

#4577 opened 2025-02-26 01:07 by copybara-service[bot]

[nnx] deepseek example v2

#4519 opened 2025-01-31 18:59 by cgarciae

[nnx] deepseek example

#4512 opened 2025-01-29 07:05 by cgarciae

deepseek example

#4511 opened 2025-01-29 07:01 by copybara-service[bot]

Improve error message for when nnx.Modules use jax or numpy arrays as leaf values

#4492 opened 2025-01-20 15:01 by RaghuSpaceRajan

[nnx] improve custom_vjp ergonomics

#4489 opened 2025-01-19 01:08 by cgarciae

[linen] Linesearch (and lbfgs) support for TrainState

#4471 opened 2025-01-06 11:36 by emiresenov

Update NNX `State.state.merge` method docs in `statelib.py`

#4453 opened 2024-12-17 00:45 by 8bitmp3

Update NNX `State.state.filter` method docs in `statelib.py`

#4452 opened 2024-12-17 00:41 by 8bitmp3

Update NNX `State.state.split` method docs in `statelib.py`

#4451 opened 2024-12-17 00:33 by 8bitmp3

Update NNX `State` class docs in `statelib.py`

#4450 opened 2024-12-17 00:25 by 8bitmp3

Lint NNX `BatchStat`, `Cache`, `Intermediate` class docs in `variablelib.py`

#4449 opened 2024-12-17 00:21 by 8bitmp3