Go
Home
Pricing
FAQ
Install
Home
Pricing
FAQ
Install
Login
via GitHub
google/flax
Pull Requests
Commits
nnx-fix-sharing
Module-split-merge
add-cached-partial
add-migration-guides
add-mypy-plugin
andsteing-patch-1
array-refs-guide
avital-notes2
avital-patch-1
avital-patch-2
avital-patch-3
avital-patch-4
avital-patch-6
box-variable
bridge-improve-bind
compact_name_scope
cudnn-lstm
cursor-2
cython-test
deepseek-example
deepseek-example-v2
dependabot/cargo/flaxlib_src/cargo-f9635f71e3
dependabot/pip/examples/imagenet/pip-22842e8f01
dependabot/pip/examples/imagenet/pip-c711b36505
deprecate-value
docs-build
docs-typo-fix-_cached_partial
dynamic-rngs
errors
eval-shape-sharding
expose-module-functions
fix-bridge-boxes
fix-checkpointing-notebook
fix-dataset-path-examples
fix-gemma
fix-jit
fix-list
fix-lm1b
fix-more-typos
fix-py314-kw-only-dataclasses
fix-site-banners
fix-summary
fix-to-linen
flax-nnx-page
flaxlib
functional-core
gemma-example-using-grain
hello-ctx
hello-ctx-my-old-friend
hijax-guide
hijax-in-graph
hijax-variable
improve-mnist-tutorial
improve-remat-docs
improve-scan
improve-scan-2
improve-tabulate-docs
keras-integration
keras-to-nnx
levskaya-patch-1
linen-fix-scan-out-axes
linen-fix-scan-out-axes-p2
linen-graph-node
linen-rtd
lsc-1743028162.4379792
main
marcvanzee-feature-request
marcvanzee-patch-1
marcvanzee-patch-2
marcvanzee-patch-3
marcvanzee-patch-4
marcvanzee-patch-5
marcvanzee-patch-6
marcvanzee-patch-7
marcvanzee-patch-8
marcvanzee-patch-9
marcvanzee-patch-10
marcvanzee-patch-11
marcvanzee-patch-12
nlp_seq_linen_mv
nlp_seq_linen_v2
nnx
nnx-add-call-to-basics
nnx-add-map-state
nnx-add-move
nnx-add-pmap
nnx-add-version
nnx-cache-flatten
nnx-custom-vjp-issue-full
nnx-debug-gc
nnx-demo
nnx-diffusion-example
nnx-empty-init
nnx-experimental
nnx-fast-flatten
nnx-fast-jit
nnx-filters-guide
nnx-fingerprint
nnx-fix-cond2
nnx-fix-jit
nnx-fix-sharing
nnx-flat-flatten
nnx-improve-filters
nnx-improve-mnist-tutorial
nnx-improve-nnx-basics
nnx-jit-new-update
nnx-leaves-are-leaves
nnx-lift-api
nnx-lm1b
nnx-mha-decode
nnx-mnist-tutorial-minor-changes
nnx-new-randomness-guide
nnx-optax-optimizer
nnx-optimizations-p4-rust
nnx-optimizations-p4
nnx-optimize-fingerprint
nnx-optimizer-use-variables
nnx-pytree-2
nnx-pytree-3
nnx-pytree-modules
nnx-pytree-state
nnx-pytree-v3
nnx-refactor-custom-vjp
nnx-refactor-graph-utils
nnx-refactor-vmap
nnx-remove-display
nnx-remove-state
nnx-remove-transform-metaclasses
nnx-runtime-defaults
nnx-simplify-flatten-unflatten
nnx-simplify-readme-again
nnx-simplify-unflatten
nnx-simplify-update
nnx-unsafe-graph-cache
nnx-update-transforms
nnx-variable-and-graph-utils
nnx-variable-issue
nnx-variable-refactor
nnx-vmap-cond
nx
pr/IvyZX/4145
pr/IvyZX/4171
pr/IvyZX/4343
pr/IvyZX/4624
pr/chiamp/3782-1
pr/chiamp/3782
pr/erfanzar/4412
pr/gnecula/3861
pr/mcsmart76/3953
quantile-regression-nnx
randomness-guide
reenable-pytest-examples
refactor-update-context
remove-variable-state
revert-626-nlp_seq_linen
rewrite-prefetch_to_device
rnn-v2
rnncell-refactor
set_mode
split-pytree
tabulate-flops-fix
tabulate-sharding
test_317378007
test_317793206
test_318181863
test_318341372
test_321144725
test_324836484
test_326589776
test_331916531
test_332047658
test_332852862
test_333036622
test_333432938
test_337964891
test_340818348
test_341637550
test_342960801
test_344339107
test_346206557
test_348184387
test_348290078
test_348290081
test_348599074
test_349699757
test_351376151
test_351584173
test_352547458
test_355497424
test_360160709
test_362256497
test_362897819
test_368110589
test_369672468
test_375880589
test_380146238
test_381300586
test_384593511
test_384833512
test_390305432
test_394394031
test_399373571
test_400667551
test_400933467
test_402700753
test_405962065
test_411882335
test_424634712
test_424833419
test_425358614
test_430675887
test_445402695
test_454835318
test_459238640
test_460812552
test_476904885
test_484569827
test_484804825
test_494218037
test_495168957
test_511585331
test_511586735
test_517661790
test_522401382
test_524020065
test_524156426
test_525967287
test_526728576
test_527795240
test_533548381
test_537384891
test_543518398
test_563486995
test_563652936
test_568217601
test_569542141
test_573301464
test_579987248
test_600454908
test_610717214
test_622872682
test_623822696
test_629615743
test_642421861
test_643090672
test_650338576
test_660085584
test_669678371
test_671493323
test_677843398
test_681582933
test_688953662
test_694151533
test_705077026
test_720657969
test_730436599
test_734367472
test_738189137
test_740051057
test_743764112
test_745698612
test_747024381
test_748311465
test_773116112
test_782979395
test_797882059
test_801930691
test_804790020
test_804851220
test_815862166
test_837235955
test_841775258
test_845027363
test_846381464
test_846723995
test-main
tf-howto
tmp
tmp2
to-linen-eval-shape
updated-pyproject-toml
use-kernel-metadata-in-gspmd-guide
use-register-dataclass
use-seq-lenghts
[nnx] don't share Rngs
cgarciae
committed
220 days ago
58fd8a1f
[nnx] add optax optimizer
cgarciae
committed
220 days ago
32fabe04
Merge pull request #4741 from google:mutable-array-p2
a-googler
committed
221 days ago
a5eebe5d
MultiHeadAttention only keeps rngs if dropout_rate is positive
Cristian Garcia
committed
222 days ago
65dfd465
move freeze, mutable, and pure to graph.py
cgarciae
committed
222 days ago
f50b4d98
Merge pull request #4749 from vfdev-5:mha-add-keep-rngs-arg
a-googler
committed
222 days ago
45ce6ee2
Merge pull request #4747 from vfdev-5:fix-docs-eval_shape
a-googler
committed
222 days ago
f265d914
Merge pull request #4746 from vfdev-5:update-pre-commit-hooks-version
a-googler
committed
222 days ago
172ac348
[flax] unconditionally register nnx.Variable as a pytree
Jake VanderPlas
committed
223 days ago
866fe952
Added keep_rngs arg to MHA to optionally store rngs
vfdev-5
committed
223 days ago
36d2421b
Fixed docstring visibility for nnx.eval_shape
vfdev-5
committed
223 days ago
e061dff5
Updated version of pre-commit-hooks in .pre-commit-config.yaml
vfdev-5
committed
224 days ago
ce4f0074
add dataclass
Cristian Garcia
committed
224 days ago
51317189
Merge pull request #4715 from google:nnx-mutable-array-p1
a-googler
committed
230 days ago
a7157cfb
[nnx] mutable array p1
cgarciae
committed
232 days ago
6ab5f83b
Update JAX nightly index usage
MichaelHudgins
committed
237 days ago
12a29ece
Merge pull request #4699 from google:nnx-jitwrapped
a-googler
committed
237 days ago
fdfc0bf4
[nnx] add JitWrapped API
cgarciae
committed
237 days ago
edc937b2
Merge pull request #4728 from google:pytree-variables
a-googler
committed
238 days ago
f3b396e7
fix merge
Cristian Garcia
committed
239 days ago
63688067
[nnx] make Variable a pytree
cgarciae
committed
243 days ago
2f3b344f
Fixes a bug in type annotations for scope.param (unbox=True should accept callable[..., T | AxisMEtadata[T]] and return T, while unbox=False should always return the same thing as what callable returning.
marksandler2
committed
243 days ago
54fad0a8
Merge pull request #4652 from jlperla:identity
a-googler
committed
244 days ago
240a5107
Merge pull request #4724 from google:0.10.6
a-googler
committed
244 days ago
0fde6365
Merge pull request #4639 from google:flaxlib-types
a-googler
committed
244 days ago
dc062922
[nnx] add more flaxlib types
cgarciae
committed
245 days ago
334f3830
Lower-bounded using new JAX release
jlperla
committed
245 days ago
5a4ec98b
Merge branch 'google:main' into identity
jlperla
committed
245 days ago
Verified
f3c11518
remove "nnx" collection in ToLinen
Cristian Garcia
committed
245 days ago
d3ec212f
v0.10.6
cgarciae
committed
245 days ago
40d7df36
Older