google/flax

Pull Requests Commits

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

815f9f53

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

cgarciae committed 1 year ago

ea57438f

Merge pull request #4191 from google:nnx-optimize-jit

a-googler committed 1 year ago

d111adf7

cgarciae committed 1 year ago

2079f15b

Merge pull request #4186 from IvyZX:errors

a-googler committed 1 year ago

b9679643

Merge pull request #4189 from IvyZX:axis-cleanup

a-googler committed 1 year ago

61ea8a66

clean up axis hooks

IvyZX committed 1 year ago

978bc3dd

Merge nnx.errors to flax.errors

IvyZX committed 1 year ago

e8b0a1fd

Merge pull request #4161 from IvyZX:bdg-logic

a-googler committed 1 year ago

f948154d

Support linen.LogicallyPartitioned <-> nnx.Variable

IvyZX committed 1 year ago

a9cb80b8

Merge pull request #4185 from google:nnx-transform-metadata-issue

a-googler committed 1 year ago

1b72435a

[nnx] enable sharding transformation on integer prefixes

cgarciae committed 1 year ago

93c06270

Merge pull request #4180 from google:nnx-fix-optimizer

a-googler committed 1 year ago

671130b3

Older