google/flax

Pull Requests Commits

Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec

yashk2810 committed 1 year ago

c1832a39

[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

[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] add cached_partial

Cristian Garcia committed 1 year ago

61e5df49

Merge pull request #4536 from mattbahr:accuracy-metric-binary-classification-support

a-googler committed 1 year ago

b5f478ff

[nnx] fix Variable hooks

Cristian Garcia committed 1 year ago

2b870fa7

[nnx] register_variable_name refactor

Cristian Garcia committed 1 year ago

c67eb094

added support to the accuracy metric for binary classification

mattbahr committed 1 year ago

08e205f8

Older