google/flax

Pull Requests Commits

Avoid usage of deprecated `jax.core` APIs.

Jake VanderPlas committed 2 hours ago

20446713

Cristian Garcia committed 22 hours ago

901efc34

use State instead of GraphState

Cristian Garcia committed 23 hours ago

09f0f930

improve check_prefix

Cristian Garcia committed 1 day ago

1c151940

Avoid deprecated JAX API

Jake VanderPlas committed 1 day ago

7ee4043d

Merge pull request #5388 from samanklesaria:rename_abstract

a-googler committed 1 day ago

d516e45a

Merge pull request #5403 from samanklesaria:diffusion

a-googler committed 1 day ago

9e73be40

Merge pull request #5370 from samanklesaria:pytree_guide_capture

a-googler committed 1 day ago

524847fb

Merge pull request #5117 from samanklesaria:opt_cookbook

a-googler committed 1 day ago

3b475c45

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

a-googler committed 1 day ago

228c50e5

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

a-googler committed 1 day ago

4c9d11fd

[JAX] Fix code that compares PartitionSpecs with tuples.

hawkinsp committed 2 days ago

dceec7d6

Add Optimization Cookbook

samanklesaria committed 2 days ago

00aa540e

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

a-googler committed 2 days ago

9bed4350

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

a-googler committed 2 days ago

4e730184

improve nnx.compat

Cristian Garcia committed 4 days ago

4fd69fa5

forward fix simple transforms by default

Cristian Garcia committed 5 days ago

b5a85d2f

Fix for change to `jax.Array.__contains__`

Jake VanderPlas committed 5 days ago

64aa65ae

Use flow matching for diffusion example

samanklesaria committed 5 days ago

0f1ef741

Avoid deprecated pattern of constructing a dtype from an array.

Jake VanderPlas committed 6 days ago

9534e266

Add diffusion example

samanklesaria committed 7 days ago

e8a62a5e

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

a-googler committed 7 days ago

a0cc8837

cgarciae committed 7 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 7 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

a-googler committed 8 days ago

d874278f

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

mohsinm-dev committed 8 days ago

4f6b33fa

Rename to 'recreate_variables'

samanklesaria committed 8 days ago

5716ff76

Add flag for variables in nnx.merge

samanklesaria committed 8 days ago

148f202b

Fix the wrong names for the hooks in nnx.with_metadata, fixed such that names match the proper hook names in Variable class.

a-googler committed 8 days ago

5a4859d2

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

danielsuo committed 8 days ago

355e141a

Older