google/flax

Pull Requests Commits

Added keep_rngs arg to MHA to optionally store rngs

vfdev-5 committed 229 days ago

36d2421b

Fixed docstring visibility for nnx.eval_shape

vfdev-5 committed 229 days ago

e061dff5

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

vfdev-5 committed 230 days ago

ce4f0074

Cristian Garcia committed 230 days ago

51317189

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

a-googler committed 236 days ago

a7157cfb

[nnx] mutable array p1

cgarciae committed 238 days ago

6ab5f83b

Update JAX nightly index usage

MichaelHudgins committed 243 days ago

12a29ece

Merge pull request #4699 from google:nnx-jitwrapped

a-googler committed 243 days ago

fdfc0bf4

[nnx] add JitWrapped API

cgarciae committed 243 days ago

edc937b2

Merge pull request #4728 from google:pytree-variables

a-googler committed 244 days ago

f3b396e7

Cristian Garcia committed 245 days ago

63688067

[nnx] make Variable a pytree

cgarciae committed 249 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 249 days ago

54fad0a8

Merge pull request #4652 from jlperla:identity

a-googler committed 250 days ago

240a5107

Merge pull request #4724 from google:0.10.6

a-googler committed 250 days ago

0fde6365

Merge pull request #4639 from google:flaxlib-types

a-googler committed 250 days ago

dc062922

[nnx] add more flaxlib types

cgarciae committed 251 days ago

334f3830

Lower-bounded using new JAX release

jlperla committed 251 days ago

5a4ec98b

Merge branch 'google:main' into identity

jlperla committed 251 days ago

Verified f3c11518

remove "nnx" collection in ToLinen

Cristian Garcia committed 251 days ago

d3ec212f

cgarciae committed 251 days ago

40d7df36

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

a-googler committed 255 days ago

dbce607b

update ml-collections

cgarciae committed 255 days ago

ccdd29f7

Make fully replicated array before saving checkpoints for examples that use pmap.

IvyZX committed 257 days ago

aa897790

Merge pull request #4630 from google:optimize-jit

a-googler committed 258 days ago

1238a616

Merge pull request #4705 from IvyZX:einsum

a-googler committed 258 days ago

63106905

Add custom einsum op to Einsum()

IvyZX committed 259 days ago

000ee89b

Rollback of attention_bias parameter, because the change overrides the attention bias for injected attention functions.

Jan Hosang committed 259 days ago

32ad6c57

Unit tests for `attention_bias` parameter to `MultiHeadDotProductAttention`. Add parameter to all overloads to make pytype happy.

Jan Hosang committed 259 days ago

5bba7b45

[nnx] refactor GraphDef

cgarciae committed 260 days ago

8cb8a7a2

Newer Older