google/flax

Pull Requests Commits

"From PyTorch to JAX and Flax" migration guide

vfdev-5 committed 86 days ago

132102fe

Addressed comments

vfdev-5 committed 86 days ago

Verified 58658acf

update version to 0.12.7

Cristian Garcia committed 87 days ago

66e58ede

Avoid usage of deprecated `jax.core` APIs.

Cristian Garcia committed 88 days ago

f6109649

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

a-googler committed 89 days ago

d8ef528a

fix 'the the' in nnx/nn/linear.py

SAY-5 committed 90 days ago

6427e72b

fix 'the the' in nnx/transforms/compilation.py

SAY-5 committed 90 days ago

3079ffd0

fix 'the the' in linen/linear.py

SAY-5 committed 90 days ago

3cb8fc94

fix 'the the' in rnglib.py

SAY-5 committed 90 days ago

884384a1

fix doubled 'the the' typos across multiple files

SAY-5 committed 90 days ago

14788d5b

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

vfdev-5 committed 91 days ago

8fa97d0c

Fixed intermedites capture

vfdev-5 committed 91 days ago

cdfadd61

Cristian Garcia committed 91 days ago

901efc34

use State instead of GraphState

Cristian Garcia committed 91 days ago

09f0f930

improve check_prefix

Cristian Garcia committed 91 days ago

1c151940

Avoid deprecated JAX API

Jake VanderPlas committed 92 days ago

7ee4043d

Merge pull request #5388 from samanklesaria:rename_abstract

a-googler committed 92 days ago

d516e45a

Merge pull request #5403 from samanklesaria:diffusion

a-googler committed 92 days ago

9e73be40

Merge pull request #5370 from samanklesaria:pytree_guide_capture

a-googler committed 92 days ago

524847fb

Merge pull request #5117 from samanklesaria:opt_cookbook

a-googler committed 92 days ago

3b475c45

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

a-googler committed 92 days ago

228c50e5

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

a-googler committed 92 days ago

4c9d11fd

[JAX] Fix code that compares PartitionSpecs with tuples.

hawkinsp committed 93 days ago

dceec7d6

Add Optimization Cookbook

samanklesaria committed 93 days ago

00aa540e

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

a-googler committed 93 days ago

9bed4350

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

a-googler committed 93 days ago

4e730184

improve nnx.compat

Cristian Garcia committed 95 days ago

4fd69fa5

forward fix simple transforms by default

Cristian Garcia committed 96 days ago

b5a85d2f

Fix for change to `jax.Array.__contains__`

Jake VanderPlas committed 96 days ago

64aa65ae

Use flow matching for diffusion example

samanklesaria committed 96 days ago

0f1ef741

Newer Older