google/flax

Pull Requests Commits

Avoid usage of deprecated `jax.core` APIs.

Jake VanderPlas committed 50 days ago

c8e3120f

Cristian Garcia committed 52 days ago

901efc34

use State instead of GraphState

Cristian Garcia committed 52 days ago

09f0f930

improve check_prefix

Cristian Garcia committed 52 days ago

1c151940

Avoid deprecated JAX API

Jake VanderPlas committed 52 days ago

7ee4043d

Merge pull request #5388 from samanklesaria:rename_abstract

a-googler committed 53 days ago

d516e45a

Merge pull request #5403 from samanklesaria:diffusion

a-googler committed 53 days ago

9e73be40

Merge pull request #5370 from samanklesaria:pytree_guide_capture

a-googler committed 53 days ago

524847fb

Merge pull request #5117 from samanklesaria:opt_cookbook

a-googler committed 53 days ago

3b475c45

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

a-googler committed 53 days ago

228c50e5

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

a-googler committed 53 days ago

4c9d11fd

[JAX] Fix code that compares PartitionSpecs with tuples.

hawkinsp committed 53 days ago

dceec7d6

Add Optimization Cookbook

samanklesaria committed 53 days ago

00aa540e

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

a-googler committed 53 days ago

9bed4350

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

a-googler committed 53 days ago

4e730184

improve nnx.compat

Cristian Garcia committed 56 days ago

4fd69fa5

forward fix simple transforms by default

Cristian Garcia committed 56 days ago

b5a85d2f

Fix for change to `jax.Array.__contains__`

Jake VanderPlas committed 56 days ago

64aa65ae

Use flow matching for diffusion example

samanklesaria committed 56 days ago

0f1ef741

Avoid deprecated pattern of constructing a dtype from an array.

Jake VanderPlas committed 57 days ago

9534e266

Add diffusion example

samanklesaria committed 58 days ago

e8a62a5e

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

a-googler committed 59 days ago

a0cc8837

cgarciae committed 59 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 59 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

a-googler committed 59 days ago

d874278f

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

mohsinm-dev committed 59 days ago

4f6b33fa

Rename to 'recreate_variables'

samanklesaria committed 59 days ago

5716ff76

Add flag for variables in nnx.merge

samanklesaria committed 59 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 59 days ago

5a4859d2

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

danielsuo committed 59 days ago

355e141a

Older