google/flax

Pull Requests Commits

[nnx] improve custom_vjp ergonomics

cgarciae committed 364 days ago

8510c161

Merge pull request #4592 from jakevdp:fix-shape

a-googler committed 1 year ago

45a8f846

Avoid calls to jnp.shape for non-array inputs.

jakevdp committed 1 year ago

4e300d4c

[bridge] improve sow and add initializers

Cristian Garcia committed 1 year ago

a24d790e

Refactor bridge.Module tests from `wrappers_test.py` to another file.

IvyZX committed 1 year ago

6af8fcb0

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

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

Older