google/flax

Pull Requests Commits

Replaced nnx.abstract_with_sharding with nnx.as_abstract

vfdev-5 committed 140 days ago

06990880

Merge pull request #5388 from samanklesaria:rename_abstract

a-googler committed 140 days ago

d516e45a

Merge pull request #5403 from samanklesaria:diffusion

a-googler committed 140 days ago

9e73be40

Merge pull request #5370 from samanklesaria:pytree_guide_capture

a-googler committed 140 days ago

524847fb

Merge pull request #5117 from samanklesaria:opt_cookbook

a-googler committed 140 days ago

3b475c45

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

a-googler committed 140 days ago

228c50e5

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

a-googler committed 140 days ago

4c9d11fd

[JAX] Fix code that compares PartitionSpecs with tuples.

hawkinsp committed 141 days ago

dceec7d6

Add Optimization Cookbook

samanklesaria committed 141 days ago

00aa540e

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

a-googler committed 141 days ago

9bed4350

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

a-googler committed 141 days ago

4e730184

improve nnx.compat

Cristian Garcia committed 143 days ago

4fd69fa5

forward fix simple transforms by default

Cristian Garcia committed 144 days ago

b5a85d2f

Fix for change to `jax.Array.__contains__`

Jake VanderPlas committed 144 days ago

64aa65ae

Use flow matching for diffusion example

samanklesaria committed 144 days ago

0f1ef741

Avoid deprecated pattern of constructing a dtype from an array.

Jake VanderPlas committed 145 days ago

9534e266

Add diffusion example

samanklesaria committed 145 days ago

e8a62a5e

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

a-googler committed 146 days ago

a0cc8837

cgarciae committed 146 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 146 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

a-googler committed 146 days ago

d874278f

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

mohsinm-dev committed 146 days ago

4f6b33fa

Rename to 'recreate_variables'

samanklesaria committed 146 days ago

5716ff76

Add flag for variables in nnx.merge

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

5a4859d2

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

danielsuo committed 147 days ago

355e141a

Merge pull request #5380 from google:removed-bonsai-from-docs

a-googler committed 148 days ago

d48b501f

improve with_rngs

Cristian Garcia committed 151 days ago

983506b3

expose nnx pooling layers in API docs

Cristian Garcia committed 151 days ago

70f4aa96

add Rngs.broadcast

Cristian Garcia committed 152 days ago

3645777b

Older