google/flax

Pull Requests Commits

migrate to pyrefly

Cristian Garcia committed 107 days ago

28c97f50

Disallow `P.UNCONSTRAINED` in `out_shardings` in `jit`. This leads to simplifications in how we lower shardings to stableHLO.

yashk2810 committed 107 days ago

9877bea7

Merge pull request #5422 from SAY-5:fix-doubled-the-typos

a-googler committed 107 days ago

09ed488a

Deprecate Variable hooks in favor of custom Variable types.

a-googler committed 107 days ago

a13993f6

update version to 0.12.7

Cristian Garcia committed 110 days ago

66e58ede

Avoid usage of deprecated `jax.core` APIs.

Cristian Garcia committed 111 days ago

f6109649

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

a-googler committed 111 days ago

d8ef528a

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

SAY-5 committed 113 days ago

6427e72b

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

SAY-5 committed 113 days ago

3079ffd0

fix 'the the' in linen/linear.py

SAY-5 committed 113 days ago

3cb8fc94

fix 'the the' in rnglib.py

SAY-5 committed 113 days ago

884384a1

fix doubled 'the the' typos across multiple files

SAY-5 committed 113 days ago

14788d5b

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

vfdev-5 committed 113 days ago

8fa97d0c

Cristian Garcia committed 114 days ago

901efc34

use State instead of GraphState

Cristian Garcia committed 114 days ago

09f0f930

improve check_prefix

Cristian Garcia committed 114 days ago

1c151940

Avoid deprecated JAX API

Jake VanderPlas committed 114 days ago

7ee4043d

Merge pull request #5388 from samanklesaria:rename_abstract

a-googler committed 115 days ago

d516e45a

Merge pull request #5403 from samanklesaria:diffusion

a-googler committed 115 days ago

9e73be40

Merge pull request #5370 from samanklesaria:pytree_guide_capture

a-googler committed 115 days ago

524847fb

Merge pull request #5117 from samanklesaria:opt_cookbook

a-googler committed 115 days ago

3b475c45

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

a-googler committed 115 days ago

228c50e5

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

a-googler committed 115 days ago

4c9d11fd

[JAX] Fix code that compares PartitionSpecs with tuples.

hawkinsp committed 115 days ago

dceec7d6

Add Optimization Cookbook

samanklesaria committed 116 days ago

00aa540e

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

a-googler committed 116 days ago

9bed4350

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

a-googler committed 116 days ago

4e730184

improve nnx.compat

Cristian Garcia committed 118 days ago

4fd69fa5

forward fix simple transforms by default

Cristian Garcia committed 118 days ago

b5a85d2f

Fix for change to `jax.Array.__contains__`

Jake VanderPlas committed 118 days ago

64aa65ae

Older