google/flax

Pull Requests Commits

use_running_average and deterministic default to None

Cristian Garcia committed 48 days ago

e3ae2c3e

Merge pull request #5421 from google:5412-add-basemod-kwargs-to-lora

a-googler committed 48 days ago

d8ef528a

Added kwargs to call to pass kwargs to basemodule in nnx.LoRA

vfdev-5 committed 50 days ago

8fa97d0c

Cristian Garcia committed 51 days ago

901efc34

use State instead of GraphState

Cristian Garcia committed 51 days ago

09f0f930

improve check_prefix

Cristian Garcia committed 51 days ago

1c151940

Avoid deprecated JAX API

Jake VanderPlas committed 51 days ago

7ee4043d

Merge pull request #5388 from samanklesaria:rename_abstract

a-googler committed 52 days ago

d516e45a

Merge pull request #5403 from samanklesaria:diffusion

a-googler committed 52 days ago

9e73be40

Merge pull request #5370 from samanklesaria:pytree_guide_capture

a-googler committed 52 days ago

524847fb

Merge pull request #5117 from samanklesaria:opt_cookbook

a-googler committed 52 days ago

3b475c45

Merge pull request #5379 from mohsinm-dev:fix-rngs-dunder-getattr

a-googler committed 52 days ago

228c50e5

Merge pull request #5381 from samanklesaria:variables-in-merge

a-googler committed 52 days ago

4c9d11fd

[JAX] Fix code that compares PartitionSpecs with tuples.

hawkinsp committed 52 days ago

dceec7d6

Add Optimization Cookbook

samanklesaria committed 52 days ago

00aa540e

Merge pull request #4913 from google:gemma-example-using-grain

a-googler committed 52 days ago

9bed4350

Merge pull request #5377 from google:fix-sow-link

a-googler committed 52 days ago

4e730184

improve nnx.compat

Cristian Garcia committed 55 days ago

4fd69fa5

forward fix simple transforms by default

Cristian Garcia committed 55 days ago

b5a85d2f

Fix for change to `jax.Array.__contains__`

Jake VanderPlas committed 55 days ago

64aa65ae

Use flow matching for diffusion example

samanklesaria committed 55 days ago

0f1ef741

Avoid deprecated pattern of constructing a dtype from an array.

Jake VanderPlas committed 56 days ago

9534e266

Add diffusion example

samanklesaria committed 57 days ago

e8a62a5e

Merge pull request #5400 from google:fix-ci-2

a-googler committed 58 days ago

a0cc8837

cgarciae committed 58 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 58 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

a-googler committed 58 days ago

d874278f

fix Rngs.__getattr__ returning default stream for dunder names (#5375)

mohsinm-dev committed 58 days ago

4f6b33fa

Rename to 'recreate_variables'

samanklesaria committed 58 days ago

5716ff76

Add flag for variables in nnx.merge

samanklesaria committed 58 days ago

148f202b

Older