google/flax

Pull Requests Commits

[nnx] add flaxlib

Cristian Garcia committed 1 year ago

dfc86d5c

Fix false positive tracer leaks in flax library.

levskaya committed 1 year ago

62a6a253

Fix error when printing module with no fields.

a-googler committed 1 year ago

c0573372

Merge pull request #4205 from google:nnx-remove-registry

a-googler committed 1 year ago

37f643da

Merge pull request #4230 from IvyZX:pure-dict

a-googler committed 1 year ago

9f44f811

Add pure dict conversion utils to .

IvyZX committed 1 year ago

3b9dce14

simplify graph.py

cgarciae committed 1 year ago

e7dba282

Merge pull request #4216 from google:nnx-randomness-guide

a-googler committed 1 year ago

8b37d1ab

[nnx] add Randomness guide

cgarciae committed 1 year ago

65d4193f

Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.

dougalm committed 1 year ago

fc19c5dc

Merge pull request #4220 from IvyZX:jit-guide

a-googler committed 1 year ago

9c86396c

Flax gspmd guide

IvyZX committed 1 year ago

9b41f61c

[nnx] revive TrainState toy example

cgarciae committed 1 year ago

2d6aa629

Merge pull request #4223 from google:fix-transforms-guide

a-googler committed 1 year ago

e3772b2d

[nnx] fix transforms guide

cgarciae committed 1 year ago

357f1665

Merge pull request #4197 from google:nnx-real-transforms-guide

a-googler committed 1 year ago

b2277abd

Forward all arguments when using nnx.transforms.deprecated.scan as a decorator.

nimrod-gileadi committed 1 year ago

2dc30a8a

fix tests for numpy 2.0 compatibility

Cristian Garcia committed 1 year ago

fcd00716

Minor documentation fixes for AxisMetadata.

a-googler committed 1 year ago

d5845516

[nnx] add transforms guide

cgarciae committed 1 year ago

fb1a9cc2

Merge pull request #4209 from IvyZX:linen-nnx-guide

a-googler committed 1 year ago

268366f1

Add trimmed Linen to NNX guide

IvyZX committed 1 year ago

fe2e5860

Merge pull request #4206 from google:fix-banners

a-googler committed 1 year ago

e356acbc

shorten banners

cgarciae committed 1 year ago

4d7c2cb7

Merge pull request #4202 from IvyZX:nnx-landing

a-googler committed 1 year ago

df5afab8

Merge pull request #4194 from IvyZX:bdg-tree

a-googler committed 1 year ago

ddaef57c

Fix landing page and banner phrasing and add examples page

IvyZX committed 1 year ago

b8fb8755

Align bridge variable tree structures

IvyZX committed 1 year ago

4a04f754

Partially revert #4192 which sets back a bunch of previous merged pushes.

IvyZX committed 1 year ago

9eb0a619

Merge pull request #4192 from google:nnx-transforms-guide

a-googler committed 1 year ago

03e034d0

Newer Older