google/flax

Pull Requests Commits

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

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

Merge pull request #4538 from google:improve-publish-safety

a-googler committed 1 year ago

33bf6b6a

update pypi publish

cgarciae committed 1 year ago

0fb09bf6

Merge pull request #4525 from google:v0.10.3

a-googler committed 1 year ago

dc01737a

Merge pull request #4524 from cool-RR:patch-5

a-googler committed 1 year ago

d935e53b

Merge pull request #4493 from google:nnx-tabulate-2

a-googler committed 1 year ago

7efdfc82

The CL introduces a new SowConfig which is used to configure which intermediate values are to be sowed.

a-googler committed 1 year ago

429033b3

cgarciae committed 1 year ago

4c1f760a

batch_norm.rst: == should be =

cool-RR committed 1 year ago

Verified cc4279c1

Merge pull request #4515 from IvyZX:perturb

a-googler committed 1 year ago

f71ce6c6

Older