google/flax

Pull Requests Commits

Merge pull request #4580 from google:nnx-tabulate-issue

a-googler committed 1 year ago

7974b106

Merge pull request #4579 from google:v0.10.4

a-googler committed 1 year ago

7f175f80

[nnx] fix tabulate

cgarciae committed 1 year ago

14fda9d7

cgarciae committed 1 year ago

5f4f1b84

[bridge] Set _initializing correctly and avoid return RNG states

IvyZX committed 1 year ago

54138501

Merge pull request #4576 from IvyZX:fix-head

a-googler committed 1 year ago

d96be6ce

Fix breakages from newest jax

IvyZX committed 1 year ago

0f9f55d6

Merge pull request #4490 from google:nnx-shard-map

a-googler committed 1 year ago

41bef07c

Add REFLECT padding to convolution layer

sarlinpe committed 1 year ago

d552d4b8

Merge pull request #4561 from IvyZX:state-dep-warn

a-googler committed 1 year ago

dd6e5955

Fix all state method usages

IvyZX committed 1 year ago

ed73ba43

[nnx] add shard_map

cgarciae committed 1 year ago

10d8e5ca

Fixes spmd to work correctly with xaot compilation by using global mesh's device instead of jax.devices()[0]

marksandler2 committed 1 year ago

1ec5ef29

Add deprecation warning to all nnx.State methods

IvyZX committed 1 year ago

14b51fd7

Merge pull request #4555 from IvyZX:linx_setup_name

a-googler committed 1 year ago

88ea291b

Add a dummy `__getattr__` to `bridge.Module` to make pytype happy in setup style.

IvyZX committed 1 year ago

013147af

Raise error if user uses 'name' in bridge module setup

IvyZX committed 1 year ago

ee377f41

Merge pull request #4554 from google:bridge-improve-stack

a-googler committed 1 year ago

f9f68854

[bridge] improve Module context

cgarciae committed 1 year ago

cb1bb173

Merge pull request #4552 from IvyZX:linx_meta

a-googler committed 1 year ago

9d2bea3e

Merge pull request #4470 from marcelroed:nnx-optimizer-typing

a-googler committed 1 year ago

7bcee15f

Add linen metadata conversion to linx

IvyZX committed 1 year ago

1c0fca00

Add `is_initializing` API

IvyZX committed 1 year ago

8d3c22b3

[bridge] fix unbox logic

Cristian Garcia committed 1 year ago

c01ff11e

Update gemma sampler to handle additional params

a-googler committed 1 year ago

d05075e1

[nnx] Pass kwargs to super `__init_subclass__`.

a-googler committed 1 year ago

323107d2

[nnx] copy _var_metadata

Cristian Garcia committed 1 year ago

eca312f1

Merge pull request #4542 from google:nnx-compact-2

a-googler committed 1 year ago

7e408079

[nnx] bridge Module

cgarciae committed 1 year ago

f2892823

[nnx] pytrees are graph nodes

cgarciae committed 1 year ago

4d8c9bbe

Newer Older