google/flax

Pull Requests Commits

Add an assert to check input ndim.

a-googler committed 2 years ago

6d2c23ad

Merge pull request #4056 from cool-RR:patch-2

a-googler committed 2 years ago

5694517f

Merge pull request #4102 from google:linen-transparent

a-googler committed 2 years ago

d20f594f

Customizing the einsums allows us to use the quantized einsum layers provided by [AQT](https://github.com/google/aqt).

a-googler committed 2 years ago

cd6218fb

set profile_duration_ms = None as in periodic_actions there's default value for both num_profile_steps and profile_duration_ms, and the profile stopping condition is when both num_profile_steps and profile_duration_ms are satisfied, so setting profile_duration_ms=None so that the passed num_profile_steps value gets used

ZhaoyueCheng committed 2 years ago

5b12e9b3

[numpy] Fix users of NumPy APIs that are removed in NumPy 2.0.

hawkinsp committed 2 years ago

5172c780

[linen] add share_scope

cgarciae committed 2 years ago

36952eaf

fix restore_checkpoint example in docstring

Cristian Garcia committed 2 years ago

46dc51eb

Merge pull request #4079 from danielwatson6:scope-shape-check

a-googler committed 2 years ago

d718655f

Merge pull request #4084 from google:nnx-split-merge-inputs

a-googler committed 2 years ago

8a880812

Merge pull request #4099 from google:nnx-reseed

a-googler committed 2 years ago

a7bdadba

Merge pull request #4081 from IvyZX:bridge

a-googler committed 2 years ago

d8bc1947

Add shape-based lazy init to LinenToNNX (prev LinenWrapper)

IvyZX committed 2 years ago

0139c90e

[nnx] add reseed

cgarciae committed 2 years ago

3f8c62fb

[nnx] add split/merge_inputs

cgarciae committed 2 years ago

82993122

[linen] fold rngs on jit to improve caching

Cristian Garcia committed 2 years ago

afcaf66d

[linen] generalize transform caching

Cristian Garcia committed 2 years ago

4e83e091

Merge pull request #4098 from google:fix-main

a-googler committed 2 years ago

65cd19e7

fix mypy failures due tu numpy update

cgarciae committed 2 years ago

30b0b0a2

Merge pull request #4092 from IvyZX:unbreak

a-googler committed 2 years ago

ceacc090

Ignore Orbax warning in deprecated flax.training and unbreak head

IvyZX committed 2 years ago

b0ebefcf

Merge pull request #4004 from google:nnx-pure

a-googler committed 2 years ago

c4066cce

Merge pull request #4059 from google:nnx-improve-filters-guide

a-googler committed 2 years ago

62127f81

[nnx] improve filters guide

cgarciae committed 2 years ago

7937ca08

Merge pull request #4077 from knightXun:patch-1

a-googler committed 2 years ago

93babfd2

Merge pull request #4078 from google:nnx-extract-apis

a-googler committed 2 years ago

e7f43dd7

[nnx] add extract APIs

cgarciae committed 2 years ago

16c2c324

Merge pull request #4076 from kaixih:add_fp8_doc

a-googler committed 2 years ago

e8021e9f

Perform shape checks for self.param AFTER unboxing, make shmap users happy

danielwatson6 committed 2 years ago

bc689087

[example]: remove lm1b useless parallism rules

knightXun committed 2 years ago

Verified 93291277

Older