google/flax

Pull Requests Commits

[bridge module] Add `bridge.share_scope` for layer-sublayer pairs.

IvyZX committed 279 days ago

8eab5753

Stop passing reduce_axes to jax.grad, jax.vjp, and jax.value_and_grad.

jburnim committed 279 days ago

f18baf07

Fix module stack typing annotation.

a-googler committed 279 days ago

9f5d1309

Merge pull request #4641 from google:nnx-improve-variable

a-googler committed 280 days ago

a9ef59b2

[nnx] improve Variable proxy for binary operations

cgarciae committed 280 days ago

6a2b33e2

Merge pull request #4634 from IvyZX:parentattr

a-googler committed 280 days ago

fa0f3e8e

Allow name arg to represent actual submodule path while keeping other attrs

IvyZX committed 282 days ago

18faf549

Merge pull request #4637 from IvyZX:custom-meta

a-googler committed 282 days ago

1e483800

allow custom axis metadata

IvyZX committed 282 days ago

a920dd22

Merge pull request #4632 from google:add-max-repr-depth

a-googler committed 282 days ago

dd3c18be

add reprlib max depth flag

cgarciae committed 283 days ago

1bdbab33

Merge pull request #4624 from IvyZX:attr-priority

a-googler committed 283 days ago

42098ad6

Add attribute path customization to bridge modules

IvyZX committed 283 days ago

c3230c02

Merge pull request #4527 from google:update-tracer-leakage-detection

a-googler committed 286 days ago

3e20cb9c

Merge pull request #4553 from sarlinpe:sarlinpe/conv-padding-reflect

a-googler committed 286 days ago

87e07186

Copybara import of the project:

Cristian Garcia committed 286 days ago

f14206cd

fix trace-level detection

cgarciae committed 286 days ago

b75de891

Merge pull request #4621 from google:nnx-fix-performance-docs

a-googler committed 287 days ago

96f9119c

fix performance guide

cgarciae committed 287 days ago

5b7b87f2

Merge pull request #4619 from IvyZX:bdg-name

a-googler committed 288 days ago

960b544f

support swap model params in gemma sampler

tianshub committed 288 days ago

bd9eddf2

Fixed typo in documentation for `save_checkpoint_multiprocess`.

agutkin committed 288 days ago

fefc4c9b

Fixed typo in `beam_search` loop.

agutkin committed 288 days ago

4741848c

allow bridge module to have 'name' field

IvyZX committed 288 days ago

6147d8db

Copybara import of the project:

a-googler committed 289 days ago

187b910e

Merge pull request #4613 from google:fix-promote-dtypes

a-googler committed 289 days ago

7c36812f

add promote_dtype as a config option for multiple layers

cgarciae committed 289 days ago

525c3d87

Merge pull request #4606 from google:nnx-standalone-variables

a-googler committed 289 days ago

fab37ad6

[nnx] add support for standalone Variables

cgarciae committed 290 days ago

d4aa2489

Merge pull request #4588 from google:nnx-graph-docs

a-googler committed 290 days ago

0362d956

Older