google/flax

Pull Requests Commits

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

a-googler committed 340 days ago

dd6e5955

Fix all state method usages

IvyZX committed 341 days ago

ed73ba43

[nnx] add shard_map

cgarciae committed 341 days ago

10d8e5ca

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

marksandler2 committed 341 days ago

1ec5ef29

Add deprecation warning to all nnx.State methods

IvyZX committed 341 days ago

14b51fd7

Merge pull request #4555 from IvyZX:linx_setup_name

a-googler committed 342 days ago

88ea291b

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

IvyZX committed 342 days ago

013147af

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

IvyZX committed 343 days ago

ee377f41

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

a-googler committed 343 days ago

f9f68854

[bridge] improve Module context

cgarciae committed 343 days ago

cb1bb173

Merge pull request #4552 from IvyZX:linx_meta

a-googler committed 343 days ago

9d2bea3e

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

a-googler committed 343 days ago

7bcee15f

Add linen metadata conversion to linx

IvyZX committed 343 days ago

1c0fca00

Add `is_initializing` API

IvyZX committed 343 days ago

8d3c22b3

[bridge] fix unbox logic

Cristian Garcia committed 344 days ago

c01ff11e

Update gemma sampler to handle additional params

a-googler committed 344 days ago

d05075e1

[nnx] Pass kwargs to super `__init_subclass__`.

a-googler committed 344 days ago

323107d2

[nnx] copy _var_metadata

Cristian Garcia committed 345 days ago

eca312f1

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

a-googler committed 347 days ago

7e408079

[nnx] bridge Module

cgarciae committed 347 days ago

f2892823

[nnx] pytrees are graph nodes

cgarciae committed 347 days ago

4d8c9bbe

[nnx] add cached_partial

Cristian Garcia committed 349 days ago

61e5df49

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

a-googler committed 349 days ago

b5f478ff

[nnx] fix Variable hooks

Cristian Garcia committed 350 days ago

2b870fa7

[nnx] register_variable_name refactor

Cristian Garcia committed 351 days ago

c67eb094

added support to the accuracy metric for binary classification

mattbahr committed 351 days ago

08e205f8

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

a-googler committed 352 days ago

33bf6b6a

update pypi publish

cgarciae committed 352 days ago

0fb09bf6

Merge pull request #4525 from google:v0.10.3

a-googler committed 352 days ago

dc01737a

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

a-googler committed 352 days ago

d935e53b

Newer Older