google/flax

Pull Requests Commits

[nnx] add optax optimizer

cgarciae committed 224 days ago

32fabe04

Merge pull request #4741 from google:mutable-array-p2

a-googler committed 225 days ago

a5eebe5d

MultiHeadAttention only keeps rngs if dropout_rate is positive

Cristian Garcia committed 226 days ago

65dfd465

move freeze, mutable, and pure to graph.py

cgarciae committed 226 days ago

f50b4d98

Merge pull request #4749 from vfdev-5:mha-add-keep-rngs-arg

a-googler committed 226 days ago

45ce6ee2

Merge pull request #4747 from vfdev-5:fix-docs-eval_shape

a-googler committed 226 days ago

f265d914

Merge pull request #4746 from vfdev-5:update-pre-commit-hooks-version

a-googler committed 226 days ago

172ac348

[flax] unconditionally register nnx.Variable as a pytree

Jake VanderPlas committed 227 days ago

866fe952

Added keep_rngs arg to MHA to optionally store rngs

vfdev-5 committed 227 days ago

36d2421b

Fixed docstring visibility for nnx.eval_shape

vfdev-5 committed 228 days ago

e061dff5

Updated version of pre-commit-hooks in .pre-commit-config.yaml

vfdev-5 committed 228 days ago

ce4f0074

Cristian Garcia committed 228 days ago

51317189

Merge pull request #4715 from google:nnx-mutable-array-p1

a-googler committed 234 days ago

a7157cfb

[nnx] mutable array p1

cgarciae committed 236 days ago

6ab5f83b

Update JAX nightly index usage

MichaelHudgins committed 241 days ago

12a29ece

Merge pull request #4699 from google:nnx-jitwrapped

a-googler committed 241 days ago

fdfc0bf4

[nnx] add JitWrapped API

cgarciae committed 241 days ago

edc937b2

Merge pull request #4728 from google:pytree-variables

a-googler committed 242 days ago

f3b396e7

Cristian Garcia committed 243 days ago

63688067

[nnx] make Variable a pytree

cgarciae committed 247 days ago

2f3b344f

Fixes a bug in type annotations for scope.param (unbox=True should accept callable[..., T | AxisMEtadata[T]] and return T, while unbox=False should always return the same thing as what callable returning.

marksandler2 committed 247 days ago

54fad0a8

Merge pull request #4652 from jlperla:identity

a-googler committed 248 days ago

240a5107

Merge pull request #4724 from google:0.10.6

a-googler committed 248 days ago

0fde6365

Merge pull request #4639 from google:flaxlib-types

a-googler committed 248 days ago

dc062922

[nnx] add more flaxlib types

cgarciae committed 249 days ago

334f3830

Lower-bounded using new JAX release

jlperla committed 249 days ago

5a4ec98b

Merge branch 'google:main' into identity

jlperla committed 249 days ago

Verified f3c11518

remove "nnx" collection in ToLinen

Cristian Garcia committed 249 days ago

d3ec212f

cgarciae committed 249 days ago

40d7df36

Merge pull request #4716 from google:update-ml-collections

a-googler committed 253 days ago

dbce607b

Older