google/flax

Pull Requests Commits

use uv package manager

cgarciae committed 1 year ago

6197477e

add uv.lock file

Cristian Garcia committed 1 year ago

cfc4989b

More testing and misc fixes on wrappers

IvyZX committed 1 year ago

a4aa7eb0

Merge pull request #4120 from dcharatan:parent_type_fix

a-googler committed 1 year ago

54634f7e

Merge pull request #4108 from google:nnx-transforms-flip

a-googler committed 1 year ago

9dba6c3f

Merge pull request #4134 from google:fix-object-state

a-googler committed 1 year ago

a6edb8ce

[nnx] fix initializing propagation

cgarciae committed 1 year ago

f6e7962d

Merge pull request #4095 from treigerm:nnx-add-groupnorm

a-googler committed 1 year ago

a3ddb0fa

Switch NNX to use Treescope instead of Penzai.

danieldjohnson committed 1 year ago

7c96b0e1

Fix link to orbax documentation

cool-RR committed 1 year ago

885d8bcf

[nnx] experimental transforms

cgarciae committed 1 year ago

faf9f67f

Merge pull request #4126 from IvyZX:bridge

a-googler committed 1 year ago

71b5a46b

IvyZX committed 1 year ago

fe1422e9

[pmap no rank reduce cleanup]: When flipping the

pschuh committed 1 year ago

e637bb39

Add logging to track deprecated codepaths.

cpgaffney1 committed 1 year ago

e6003511

avoid mixing `einsum_dot_general` and `einsum` argument by specifying them explicitly in the caller.

a-googler committed 1 year ago

8e0228e7

Fix `_ParentType` annotation

dcharatan committed 1 year ago

Verified 967353f6

Add GroupNorm to NNX normalization layers

treigerm committed 1 year ago

b0236062

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

a-googler committed 1 year ago

5694517f

Merge pull request #4102 from google:linen-transparent

a-googler committed 1 year ago

d20f594f

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

a-googler committed 1 year 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 1 year ago

5b12e9b3

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

hawkinsp committed 1 year ago

5172c780

add JAX-style NNX Transforms FLIP

cgarciae committed 1 year ago

da5192d0

[linen] add share_scope

cgarciae committed 1 year ago

36952eaf

fix restore_checkpoint example in docstring

Cristian Garcia committed 1 year ago

46dc51eb

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

a-googler committed 1 year ago

d718655f

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

a-googler committed 1 year ago

8a880812

Merge pull request #4099 from google:nnx-reseed

a-googler committed 1 year ago

a7bdadba

Merge pull request #4081 from IvyZX:bridge

a-googler committed 1 year ago

d8bc1947

Newer Older