google/flax

Pull Requests Commits

[nnx] improve custom_vjp ergonomics

cgarciae committed 296 days ago

8510c161

Merge pull request #4592 from jakevdp:fix-shape

a-googler committed 297 days ago

45a8f846

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

jakevdp committed 298 days ago

4e300d4c

[bridge] improve sow and add initializers

Cristian Garcia committed 301 days ago

a24d790e

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

IvyZX committed 303 days ago

6af8fcb0

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

a-googler committed 303 days ago

7974b106

Merge pull request #4579 from google:v0.10.4

a-googler committed 303 days ago

7f175f80

[nnx] fix tabulate

cgarciae committed 303 days ago

14fda9d7

cgarciae committed 303 days ago

5f4f1b84

[bridge] Set _initializing correctly and avoid return RNG states

IvyZX committed 304 days ago

54138501

Merge pull request #4576 from IvyZX:fix-head

a-googler committed 304 days ago

d96be6ce

Fix breakages from newest jax

IvyZX committed 304 days ago

0f9f55d6

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

a-googler committed 305 days ago

41bef07c

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

a-googler committed 308 days ago

dd6e5955

Fix all state method usages

IvyZX committed 308 days ago

ed73ba43

[nnx] add shard_map

cgarciae committed 308 days ago

10d8e5ca

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

marksandler2 committed 309 days ago

1ec5ef29

Add deprecation warning to all nnx.State methods

IvyZX committed 309 days ago

14b51fd7

Merge pull request #4555 from IvyZX:linx_setup_name

a-googler committed 309 days ago

88ea291b

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

IvyZX committed 309 days ago

013147af

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

IvyZX committed 310 days ago

ee377f41

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

a-googler committed 310 days ago

f9f68854

[bridge] improve Module context

cgarciae committed 310 days ago

cb1bb173

Merge pull request #4552 from IvyZX:linx_meta

a-googler committed 310 days ago

9d2bea3e

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

a-googler committed 311 days ago

7bcee15f

Add linen metadata conversion to linx

IvyZX committed 311 days ago

1c0fca00

Add `is_initializing` API

IvyZX committed 311 days ago

8d3c22b3

[bridge] fix unbox logic

Cristian Garcia committed 311 days ago

c01ff11e

Update gemma sampler to handle additional params

a-googler committed 311 days ago

d05075e1

[nnx] Pass kwargs to super `__init_subclass__`.

a-googler committed 311 days ago

323107d2

Older