google/flax

Pull Requests Commits

Replaced nnx.abstract_with_sharding with nnx.as_abstract

vfdev-5 committed 117 days ago

06990880

Merge pull request #5388 from samanklesaria:rename_abstract

a-googler committed 118 days ago

d516e45a

Merge pull request #5403 from samanklesaria:diffusion

a-googler committed 118 days ago

9e73be40

Merge pull request #5370 from samanklesaria:pytree_guide_capture

a-googler committed 118 days ago

524847fb

Merge pull request #5117 from samanklesaria:opt_cookbook

a-googler committed 118 days ago

3b475c45

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

a-googler committed 118 days ago

228c50e5

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

a-googler committed 118 days ago

4c9d11fd

[JAX] Fix code that compares PartitionSpecs with tuples.

hawkinsp committed 118 days ago

dceec7d6

Add Optimization Cookbook

samanklesaria committed 118 days ago

00aa540e

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

a-googler committed 118 days ago

9bed4350

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

a-googler committed 118 days ago

4e730184

improve nnx.compat

Cristian Garcia committed 121 days ago

4fd69fa5

forward fix simple transforms by default

Cristian Garcia committed 121 days ago

b5a85d2f

Fix for change to `jax.Array.__contains__`

Jake VanderPlas committed 121 days ago

64aa65ae

Use flow matching for diffusion example

samanklesaria committed 121 days ago

0f1ef741

Avoid deprecated pattern of constructing a dtype from an array.

Jake VanderPlas committed 122 days ago

9534e266

Add diffusion example

samanklesaria committed 123 days ago

e8a62a5e

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

a-googler committed 124 days ago

a0cc8837

cgarciae committed 124 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 124 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

a-googler committed 124 days ago

d874278f

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

mohsinm-dev committed 124 days ago

4f6b33fa

Rename to 'recreate_variables'

samanklesaria committed 124 days ago

5716ff76

Add flag for variables in nnx.merge

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

5a4859d2

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

danielsuo committed 124 days ago

355e141a

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

a-googler committed 125 days ago

d48b501f

improve with_rngs

Cristian Garcia committed 129 days ago

983506b3

expose nnx pooling layers in API docs

Cristian Garcia committed 129 days ago

70f4aa96

add Rngs.broadcast

Cristian Garcia committed 129 days ago

3645777b

Older