https://github.com/google/jax
- HEAD
- refs/heads/LenaMartens-patch-1
- refs/heads/ad-on-fancy-types
- refs/heads/add-nnx-reference
- refs/heads/add-tpu-core-count
- refs/heads/adx3
- refs/heads/adx4
- refs/heads/apaszke-debug-gpu
- refs/heads/array_tutorial
- refs/heads/avals-with-names
- refs/heads/bazel_tpu
- refs/heads/cache_log
- refs/heads/calltf_test
- refs/heads/ci_v3-8
- refs/heads/cse
- refs/heads/debug-nans-test-block-async
- refs/heads/dependabot/github_actions/actions/setup-python-5.2.0
- refs/heads/dependabot/github_actions/actions/upload-artifact-4.4.0
- refs/heads/dependabot/pip/contourpy-1.3.0
- refs/heads/dependabot/pip/etils-epathepy--1.9.3
- refs/heads/dependabot/pip/hypothesis-6.111.2
- refs/heads/dependabot/pip/rich-13.8.0
- refs/heads/dependabot/pip/setuptools-74.0.0
- refs/heads/disable_test
- refs/heads/doc_update
- refs/heads/dynamic-experiments
- refs/heads/dynamic-scoping
- refs/heads/effect-types
- refs/heads/fix
- refs/heads/fix-notebooks
- refs/heads/fix-notebooks2
- refs/heads/fix_cloud_tpu_check
- refs/heads/fixed-point
- refs/heads/furo
- refs/heads/gda1
- refs/heads/gda2
- refs/heads/gmap
- refs/heads/gnecula-patch-1
- refs/heads/gpu-determinism-note
- refs/heads/hmm-example
- refs/heads/hoist-consts
- refs/heads/ijit
- refs/heads/index
- refs/heads/initial-state
- refs/heads/initial-style-autodidax
- refs/heads/issue2263
- refs/heads/issue3040
- refs/heads/issue3285
- refs/heads/issue3620
- refs/heads/issue768
- refs/heads/jax-attrs
- refs/heads/jax_release_0.3.6
- refs/heads/jb/broadcast-tie-in
- refs/heads/jb/debug-strong-zeros
- refs/heads/jb/partial-eval-scan
- refs/heads/jb/refs
- refs/heads/jb/scan-consts
- refs/heads/jb/scan-consts-wip
- refs/heads/jb/staxperiments-wip
- refs/heads/jb/tagging
- refs/heads/jet
- refs/heads/jet2
- refs/heads/kokoro-fix
- refs/heads/layer-scan-remat
- refs/heads/libtpu_import_fix
- refs/heads/libtpu_install
- refs/heads/linearize
- refs/heads/lint-action-status
- refs/heads/log-resharding
- refs/heads/main
- refs/heads/make-custom-vjp-bwd-nones-more-robust
- refs/heads/masking-revisions
- refs/heads/maxfail
- refs/heads/mlir
- refs/heads/mutable-array-scan
- refs/heads/nanobind
- refs/heads/nightly
- refs/heads/no-more-linear-util-temp
- refs/heads/no-more-post-process
- refs/heads/on_demand_ci_jax_gpu
- refs/heads/pargmax-primitive
- refs/heads/pdot
- refs/heads/pjrt_c_api
- refs/heads/pjrt_c_api_tests
- refs/heads/pjrt_tpu
- refs/heads/pmap
- refs/heads/precision-flag-name-updating
- refs/heads/pretty-print-improvements
- refs/heads/primitives
- refs/heads/print-partial-eval-stats
- refs/heads/print_refactor
- refs/heads/prngkey-linearity-brainstorming
- refs/heads/profiler_test
- refs/heads/pullrequest
- refs/heads/pytree-attrs
- refs/heads/rafi
- refs/heads/random-docstring-fix
- refs/heads/ray
- refs/heads/refs-in-vjps
- refs/heads/rejames3
- refs/heads/restore_jax_remat_opt_barrier_4.2.0
- refs/heads/revive-leak-checker
- refs/heads/rvjp
- refs/heads/sc
- refs/heads/setupcfg
- refs/heads/sharadmv-patch-1
- refs/heads/sharadmv-patch-2
- refs/heads/sharding_docstrings
- refs/heads/shoyer-jit-stuff
- refs/heads/side-effects
- refs/heads/skye-patch-1
- refs/heads/skye-patch-2
- refs/heads/skye-patch-3
- refs/heads/skye-patch-4
- refs/heads/skye_tpu_ci_testing
- refs/heads/software-pipeline
- refs/heads/solvers
- refs/heads/source-line-info-experiments
- refs/heads/speed-up-hips
- refs/heads/stackless
- refs/heads/stax2k
- refs/heads/staxperiments
- refs/heads/test_336130333
- refs/heads/test_351557292
- refs/heads/test_355704839
- refs/heads/test_357720302
- refs/heads/test_363455815
- refs/heads/test_372280020
- refs/heads/test_377534845
- refs/heads/test_381553014
- refs/heads/test_388304561
- refs/heads/test_388307100
- refs/heads/test_390018950
- refs/heads/test_394110078
- refs/heads/test_395690233
- refs/heads/test_397174680
- refs/heads/test_398845683
- refs/heads/test_401921218
- refs/heads/test_404746466
- refs/heads/test_407577871
- refs/heads/test_409589575
- refs/heads/test_410618045
- refs/heads/test_410822358
- refs/heads/test_412110232
- refs/heads/test_415522599
- refs/heads/test_417636490
- refs/heads/test_418424182
- refs/heads/test_419001739
- refs/heads/test_419670164
- refs/heads/test_420040086
- refs/heads/test_423358730
- refs/heads/test_423914979
- refs/heads/test_424101639
- refs/heads/test_428105412
- refs/heads/test_428681711
- refs/heads/test_429097675
- refs/heads/test_429404946
- refs/heads/test_429407854
- refs/heads/test_430363082
- refs/heads/test_431602697
- refs/heads/test_432560862
- refs/heads/test_433530047
- refs/heads/test_435398406
- refs/heads/test_435759524
- refs/heads/test_436391492
- refs/heads/test_437633657
- refs/heads/test_438033664
- refs/heads/test_440450701
- refs/heads/test_440927765
- refs/heads/test_441041293
- refs/heads/test_442162369
- refs/heads/test_442573037
- refs/heads/test_444674770
- refs/heads/test_445478827
- refs/heads/test_446530978
- refs/heads/test_447366636
- refs/heads/test_448898047
- refs/heads/test_449567602
- refs/heads/test_452087344
- refs/heads/test_454250510
- refs/heads/test_454708810
- refs/heads/test_456562272
- refs/heads/test_456931912
- refs/heads/test_457007257
- refs/heads/test_458532437
- refs/heads/test_459875360
- refs/heads/test_460144999
- refs/heads/test_460216807
- refs/heads/test_460531663
- refs/heads/test_460834259
- refs/heads/test_461075508
- refs/heads/test_461210795
- refs/heads/test_462219572
- refs/heads/test_464830257
- refs/heads/test_464993276
- refs/heads/test_465655514
- refs/heads/test_466488206
- refs/heads/test_467774521
- refs/heads/test_468028234
- refs/heads/test_469397862
- refs/heads/test_470094033
- refs/heads/test_471352829
- refs/heads/test_471795639
- refs/heads/test_472517963
- refs/heads/test_477322238
- refs/heads/test_478456770
- refs/heads/test_480202723
- refs/heads/test_480224686
- refs/heads/test_482313959
- refs/heads/test_482602398
- refs/heads/test_482900882
- refs/heads/test_482900887
- refs/heads/test_482900889
- refs/heads/test_482900892
- refs/heads/test_482966582
- refs/heads/test_484448150
- refs/heads/test_487929724
- refs/heads/test_489763439
- refs/heads/test_490629581
- refs/heads/test_490662809
- refs/heads/test_493164942
- refs/heads/test_493596333
- refs/heads/test_493952215
- refs/heads/test_496483150
- refs/heads/test_496643067
- refs/heads/test_497091885
- refs/heads/test_497225669
- refs/heads/test_497604432
- refs/heads/test_497854051
- refs/heads/test_498025158
- refs/heads/test_498264649
- refs/heads/test_498270669
- refs/heads/test_498435892
- refs/heads/test_500618944
- refs/heads/test_501135700
- refs/heads/test_501345828
- refs/heads/test_502156425
- refs/heads/test_502769787
- refs/heads/test_502853874
- refs/heads/test_503474380
- refs/heads/test_505159976
- refs/heads/test_505998738
- refs/heads/test_507395004
- refs/heads/test_508168519
- refs/heads/test_508168521
- refs/heads/test_508168522
- refs/heads/test_508464286
- refs/heads/test_508726173
- refs/heads/test_509261214
- refs/heads/test_509637208
- refs/heads/test_512183375
- refs/heads/test_512192439
- refs/heads/test_514760173
- refs/heads/test_515135429
- refs/heads/test_515216183
- refs/heads/test_515494408
- refs/heads/test_516810517
- refs/heads/test_516810518
- refs/heads/test_516810520
- refs/heads/test_516810523
- refs/heads/test_517283818
- refs/heads/test_517772432
- refs/heads/test_518102665
- refs/heads/test_518506421
- refs/heads/test_518850240
- refs/heads/test_518859003
- refs/heads/test_518870432
- refs/heads/test_520744762
- refs/heads/test_522180444
- refs/heads/test_522435255
- refs/heads/test_523544053
- refs/heads/test_526822983
- refs/heads/test_527801553
- refs/heads/test_530376469
- refs/heads/test_530618498
- refs/heads/test_533228031
- refs/heads/test_536873968
- refs/heads/test_539130561
- refs/heads/test_541092273
- refs/heads/test_541402440
- refs/heads/test_541402556
- refs/heads/test_542362963
- refs/heads/test_545420249
- refs/heads/test_545655471
- refs/heads/test_547626617
- refs/heads/test_547626618
- refs/heads/test_547906057
- refs/heads/test_547924761
- refs/heads/test_548696078
- refs/heads/test_548696079
- refs/heads/test_549523682
- refs/heads/test_551644794
- refs/heads/test_552737402
- refs/heads/test_552949448
- refs/heads/test_552950630
- refs/heads/test_553240003
- refs/heads/test_553454659
- refs/heads/test_554489856
- refs/heads/test_554506940
- refs/heads/test_554515251
- refs/heads/test_555009598
- refs/heads/test_555362945
- refs/heads/test_557405952
- refs/heads/test_557606593
- refs/heads/test_557639059
- refs/heads/test_559223667
- refs/heads/test_559235820
- refs/heads/test_559282859
- refs/heads/test_559537321
- refs/heads/test_559832084
- refs/heads/test_560102989
- refs/heads/test_562875610
- refs/heads/test_563106152
- refs/heads/test_563852473
- refs/heads/test_564778874
- refs/heads/test_564832306
- refs/heads/test_565394366
- refs/heads/test_565598091
- refs/heads/test_566338492
- refs/heads/test_566829535
- refs/heads/test_567743154
- refs/heads/test_567883345
- refs/heads/test_568984719
- refs/heads/test_569230764
- refs/heads/test_570062727
- refs/heads/test_570610183
- refs/heads/test_573316975
- refs/heads/test_573862211
- refs/heads/test_573943286
- refs/heads/test_575248203
- refs/heads/test_577461132
- refs/heads/test_578841721
- refs/heads/test_578847678
- refs/heads/test_578946432
- refs/heads/test_579940420
- refs/heads/test_581954801
- refs/heads/test_584888085
- refs/heads/test_587518619
- refs/heads/test_587806540
- refs/heads/test_588037215
- refs/heads/test_588041819
- refs/heads/test_589017719
- refs/heads/test_589018175
- refs/heads/test_589025336
- refs/heads/test_591066347
- refs/heads/test_592754615
- refs/heads/test_592998231
- refs/heads/test_595845422
- refs/heads/test_595886400
- refs/heads/test_595897224
- refs/heads/test_597341374
- refs/heads/test_597350925
- refs/heads/test_597796534
- refs/heads/test_598509568
- refs/heads/test_606813605
- refs/heads/test_607616480
- refs/heads/test_607748347
- refs/heads/test_607752766
- refs/heads/test_607900615
- refs/heads/test_608638040
- refs/heads/test_610593276
- refs/heads/test_611694090
- refs/heads/test_611798938
- refs/heads/test_613248331
- refs/heads/test_613354286
- refs/heads/test_613369789
- refs/heads/test_613418799
- refs/heads/test_613719003
- refs/heads/test_614048400
- refs/heads/test_614361032
- refs/heads/test_614684199
- refs/heads/test_616247522
- refs/heads/test_616869664
- refs/heads/test_617862518
- refs/heads/test_618058831
- refs/heads/test_618721726
- refs/heads/test_618824221
- refs/heads/test_619414418
- refs/heads/test_619787143
- refs/heads/test_620111786
- refs/heads/test_620893492
- refs/heads/test_622817772
- refs/heads/test_622843729
- refs/heads/test_627162247
- refs/heads/test_629530367
- refs/heads/test_630342579
- refs/heads/test_631169957
- refs/heads/test_631422730
- refs/heads/test_632221941
- refs/heads/test_632679634
- refs/heads/test_634465693
- refs/heads/test_634604111
- refs/heads/test_636160192
- refs/heads/test_636229966
- refs/heads/test_636275268
- refs/heads/test_636578063
- refs/heads/test_636911659
- refs/heads/test_636912806
- refs/heads/test_636938648
- refs/heads/test_638687267
- refs/heads/test_638835261
- refs/heads/test_638914760
- refs/heads/test_639262876
- refs/heads/test_639830807
- refs/heads/test_639919243
- refs/heads/test_639930231
- refs/heads/test_640471925
- refs/heads/test_640500925
- refs/heads/test_640531895
- refs/heads/test_640730760
- refs/heads/test_640863826
- refs/heads/test_640960141
- refs/heads/test_641165459
- refs/heads/test_641313727
- refs/heads/test_642008162
- refs/heads/test_643201055
- refs/heads/test_643236017
- refs/heads/test_643244202
- refs/heads/test_643443500
- refs/heads/test_643565877
- refs/heads/test_643801849
- refs/heads/test_644909257
- refs/heads/test_645167590
- refs/heads/test_645227183
- refs/heads/test_645495562
- refs/heads/test_646996565
- refs/heads/test_647647825
- refs/heads/test_648738847
- refs/heads/test_648766793
- refs/heads/test_649054291
- refs/heads/test_649062595
- refs/heads/test_649089639
- refs/heads/test_650294623
- refs/heads/test_651125782
- refs/heads/test_651194020
- refs/heads/test_651464925
- refs/heads/test_651670663
- refs/heads/test_651958321
- refs/heads/test_653982391
- refs/heads/test_656180397
- refs/heads/test_657288382
- refs/heads/test_658023765
- refs/heads/test_658208973
- refs/heads/test_659497280
- refs/heads/test_659506512
- refs/heads/test_659680324
- refs/heads/test_659894314
- refs/heads/test_660355883
- refs/heads/test_660808213
- refs/heads/test_660816742
- refs/heads/test_662237972
- refs/heads/test_662484187
- refs/heads/test_663415374
- refs/heads/test_665103388
- refs/heads/test_665150506
- refs/heads/test_665741391
- refs/heads/test_666323643
- refs/heads/test_666557321
- refs/heads/test_667576309
- refs/heads/test_667665618
- refs/heads/test_667933948
- refs/heads/test_668170852
- refs/heads/test_668174706
- refs/heads/test_668538663
- refs/heads/test_668990280
- refs/heads/test_669053167
- refs/heads/test_669109443
- refs/heads/test_669129415
- refs/heads/test_669138292
- refs/heads/test_669243043
- refs/heads/test_669307421
- refs/heads/test_669307422
- refs/heads/test_669307423
- refs/heads/test_669344855
- refs/heads/test_669347182
- refs/heads/test_670053348
- refs/heads/test_670205589
- refs/heads/test_670382729
- refs/heads/test_670657245
- refs/heads/test_670672823
- refs/heads/test_670913460
- refs/heads/test_670925669
- refs/heads/test_671000480
- refs/heads/test_671707572
- refs/heads/test_672304526
- refs/heads/test_672310292
- refs/heads/test_fix
- refs/heads/test_no_log_spam
- refs/heads/test_timeout_tpu
- refs/heads/timeout
- refs/heads/tpu-ext-macos
- refs/heads/tpu_build_from_head
- refs/heads/tpu_ci_disable_tcmalloc
- refs/heads/tpu_ci_fix_compat_build
- refs/heads/tpu_ci_pjrt
- refs/heads/tpu_ci_remove_workaround
- refs/heads/tpu_ci_strings
- refs/heads/tpu_py_version
- refs/heads/tpu_py_version2
- refs/heads/tweaks
- refs/heads/typo
- refs/heads/undo-tree
- refs/heads/update-pypi
- refs/heads/variance_scaling_axes
- refs/heads/yashk2810-patch-1
- refs/heads/yashk2810-patch-10
- refs/heads/yashk2810-patch-11
- refs/heads/yashk2810-patch-12
- refs/heads/yashk2810-patch-13
- refs/heads/yashk2810-patch-14
- refs/heads/yashk2810-patch-15
- refs/heads/yashk2810-patch-16
- refs/heads/yashk2810-patch-17
- refs/heads/yashk2810-patch-18
- refs/heads/yashk2810-patch-19
- refs/heads/yashk2810-patch-2
- refs/heads/yashk2810-patch-20
- refs/heads/yashk2810-patch-3
- refs/heads/yashk2810-patch-4
- refs/heads/yashk2810-patch-5
- refs/heads/yashk2810-patch-6
- refs/heads/yashk2810-patch-7
- refs/heads/yashk2810-patch-8
- refs/heads/yashk2810-patch-9
- refs/remotes/upstream/rename-mastertrace
- refs/tags/0.3.20
- refs/tags/0.3.21
- refs/tags/0.3.22
- refs/tags/0.3.23
- refs/tags/0.3.24
- refs/tags/0.3.25
- refs/tags/jax-v0.1.49
- refs/tags/jax-v0.1.55
- refs/tags/jax-v0.1.58
- refs/tags/jax-v0.1.70
- refs/tags/jax-v0.2.20
- refs/tags/jax-v0.2.21
- refs/tags/jax-v0.2.23
- refs/tags/jax-v0.2.24
- refs/tags/jax-v0.2.26
- refs/tags/jax-v0.3.0
- refs/tags/jax-v0.3.1
- refs/tags/jax-v0.3.10
- refs/tags/jax-v0.3.11
- refs/tags/jax-v0.3.12
- refs/tags/jax-v0.3.21-rc
- refs/tags/jax-v0.3.6
- refs/tags/jax-v0.3.9
- refs/tags/jax-v0.4.29-rc
- refs/tags/jaxlib-v0.1.32
- refs/tags/jaxlib-v0.1.33
- refs/tags/jaxlib-v0.1.34
- refs/tags/jaxlib-v0.1.35
- refs/tags/jaxlib-v0.1.36
- refs/tags/jaxlib-v0.1.37
- refs/tags/jaxlib-v0.1.38
- refs/tags/jaxlib-v0.1.39
- refs/tags/jaxlib-v0.1.40
- refs/tags/jaxlib-v0.1.42
- refs/tags/jaxlib-v0.1.43
- refs/tags/jaxlib-v0.1.44
- refs/tags/jaxlib-v0.1.45
- refs/tags/jaxlib-v0.1.46
- refs/tags/jaxlib-v0.1.47
- refs/tags/jaxlib-v0.1.50
- refs/tags/jaxlib-v0.1.51
- refs/tags/jaxlib-v0.1.55
- refs/tags/jaxlib-v0.1.70
- refs/tags/jaxlib-v0.1.71
- refs/tags/jaxlib-v0.1.72
- refs/tags/jaxlib-v0.1.73
- refs/tags/jaxlib-v0.1.75
- refs/tags/jaxlib-v0.1.76
- refs/tags/jaxlib-v0.3.0
- refs/tags/jaxlib-v0.3.10
- refs/tags/test-jaxlib
- refs/tags/workflow_test
- test-docs
- jaxlib-v0.4.9
- jaxlib-v0.4.7
- jaxlib-v0.4.6
- jaxlib-v0.4.4
- jaxlib-v0.4.31
- jaxlib-v0.4.30
- jaxlib-v0.4.3
- jaxlib-v0.4.29
- jaxlib-v0.4.28
- jaxlib-v0.4.27
- jaxlib-v0.4.26
- jaxlib-v0.4.25
- jaxlib-v0.4.24
- jaxlib-v0.4.23
- jaxlib-v0.4.22
- jaxlib-v0.4.21
- jaxlib-v0.4.20
- jaxlib-v0.4.2
- jaxlib-v0.4.19
- jaxlib-v0.4.18
- jaxlib-v0.4.17
- jaxlib-v0.4.16
- jaxlib-v0.4.15
- jaxlib-v0.4.14
- jaxlib-v0.4.13
- jaxlib-v0.4.12
- jaxlib-v0.4.11
- jaxlib-v0.4.10
- jaxlib-v0.4.1
- jaxlib-v0.4.0
- jaxlib-v0.3.7
- jaxlib-v0.3.5
- jaxlib-v0.3.25
- jaxlib-v0.3.24
- jaxlib-v0.3.22
- jaxlib-v0.3.20
- jaxlib-v0.3.2
- jaxlib-v0.3.15
- jaxlib-v0.3.14
- jaxlib-v0.1.74
- jaxlib-v0.1.69
- jaxlib-v0.1.68
- jaxlib-v0.1.67
- jaxlib-v0.1.66
- jaxlib-v0.1.65
- jaxlib-v0.1.64
- jaxlib-v0.1.63
- jaxlib-v0.1.62
- jaxlib-v0.1.61
- jaxlib-v0.1.60
- jaxlib-v0.1.59
- jaxlib-v0.1.58
- jaxlib-v0.1.57
- jaxlib-v0.1.56
- jaxlib-v0.1.52
- jaxlib-v0.1.49
- jaxlib-v0.1.48
- jax-v0.4.9-rc
- jax-v0.4.9
- jax-v0.4.8
- jax-v0.4.7-rc1
- jax-v0.4.7-rc
- jax-v0.4.7
- jax-v0.4.6-rc
- jax-v0.4.6
- jax-v0.4.5
- jax-v0.4.4-rc
- jax-v0.4.4
- jax-v0.4.31-rc
- jax-v0.4.31
- jax-v0.4.30-rc
- jax-v0.4.30
- jax-v0.4.3-rc
- jax-v0.4.3
- jax-v0.4.29
- jax-v0.4.28-rc
- jax-v0.4.28
- jax-v0.4.27-rc
- jax-v0.4.27
- jax-v0.4.26-rc
- jax-v0.4.26
- jax-v0.4.25-rc
- jax-v0.4.25
- jax-v0.4.24-rc
- jax-v0.4.24
- jax-v0.4.23-rc
- jax-v0.4.23
- jax-v0.4.22-rc2
- jax-v0.4.22-rc
- jax-v0.4.22
- jax-v0.4.21-rc
- jax-v0.4.21
- jax-v0.4.20-rc
- jax-v0.4.20
- jax-v0.4.2-rc
- jax-v0.4.2
- jax-v0.4.19-rc
- jax-v0.4.19
- jax-v0.4.18-rc
- jax-v0.4.18
- jax-v0.4.17-rc
- jax-v0.4.17
- jax-v0.4.16-rc1
- jax-v0.4.16-rc
- jax-v0.4.16
- jax-v0.4.15-rc
- jax-v0.4.15
- jax-v0.4.14-rc
- jax-v0.4.14
- jax-v0.4.13-rc
- jax-v0.4.13
- jax-v0.4.12-rc
- jax-v0.4.12
- jax-v0.4.11-rc
- jax-v0.4.11
- jax-v0.4.10-rc
- jax-v0.4.10
- jax-v0.4.1-rc
- jax-v0.4.1
- jax-v0.4.0-rc1
- jax-v0.4.0-rc
- jax-v0.4.0
- jax-v0.3.8
- jax-v0.3.7
- jax-v0.3.5
- jax-v0.3.4
- jax-v0.3.3
- jax-v0.3.25-rc3
- jax-v0.3.25-rc2
- jax-v0.3.25-rc1
- jax-v0.3.25-rc
- jax-v0.3.25
- jax-v0.3.24-rc1
- jax-v0.3.24-rc
- jax-v0.3.24
- jax-v0.3.23-rc
- jax-v0.3.23
- jax-v0.3.22-rc
- jax-v0.3.22
- jax-v0.3.21
- jax-v0.3.20-rc
- jax-v0.3.20
- jax-v0.3.2
- jax-v0.3.19
- jax-v0.3.18-rc
- jax-v0.3.18
- jax-v0.3.17
- jax-v0.3.16
- jax-v0.3.15-rc
- jax-v0.3.15
- jax-v0.3.14
- jax-v0.3.13
- jax-v0.2.9
- jax-v0.2.8
- jax-v0.2.7
- jax-v0.2.6
- jax-v0.2.5
- jax-v0.2.4
- jax-v0.2.3
- jax-v0.2.28
- jax-v0.2.27
- jax-v0.2.25
- jax-v0.2.22
- jax-v0.2.2
- jax-v0.2.19
- jax-v0.2.18
- jax-v0.2.17
- jax-v0.2.16
- jax-v0.2.15
- jax-v0.2.14
- jax-v0.2.13
- jax-v0.2.12
- jax-v0.2.11
- jax-v0.2.10
- jax-v0.2.1
- jax-v0.2.0
- jax-v0.1.77
- jax-v0.1.76
- jax-v0.1.75
- jax-v0.1.74
- jax-v0.1.73
- jax-v0.1.72
- jax-v0.1.71
- jax-v0.1.69
- jax-v0.1.68
- jax-v0.1.67
- jax-v0.1.66
- jax-v0.1.65
- jax-v0.1.64
- jax-v0.1.63
- jax-v0.1.62
- jax-v0.1.61
- jax-v0.1.60
- jax-v0.1.59
Raw File
Take a new snapshot of a software origin
If the archived software origin currently browsed is not synchronized with its upstream version (for instance when new commits have been issued), you can explicitly request Software Heritage to take a new snapshot of it.
Use the form below to proceed. Once a request has been submitted and accepted, it will be processed as soon as possible. You can then check its processing state by visiting this dedicated page.Processing "take a new snapshot" request ...
Permalinks
To reference or cite the objects present in the Software Heritage archive, permalinks based on SoftWare Hash IDentifiers (SWHIDs) must be used.
Select below a type of object currently browsed in order to display its associated SWHID and permalink.
Tip revision: cfa94899bf56fc2c763ac8e05630ac960ce39264 authored by Tomás Longeri on 10 July 2024, 23:46:53 UTC
[Mosaic] Add "Not implemented" prefix to error message.
[Mosaic] Add "Not implemented" prefix to error message.
Tip revision: cfa9489
typing.py
# Copyright 2023 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
The JAX typing module is where JAX-specific static type annotations live.
This submodule is a work in progress; to see the proposal behind the types exported
here, see https://jax.readthedocs.io/en/latest/jep/12049-type-annotations.html.
The currently-available types are:
- :class:`jax.Array`: annotation for any JAX array or tracer (i.e. representations of arrays
within JAX transforms).
- :obj:`jax.typing.ArrayLike`: annotation for any value that is safe to implicitly cast to
a JAX array; this includes :class:`jax.Array`, :class:`numpy.ndarray`, as well as Python
builtin numeric values (e.g. :class:`int`, :class:`float`, etc.) and numpy scalar values
(e.g. :class:`numpy.int32`, :class:`numpy.flota64`, etc.)
- :obj:`jax.typing.DTypeLike`: annotation for any value that can be cast to a JAX-compatible
dtype; this includes strings (e.g. `'float32'`, `'int32'`), scalar types (e.g. `float`,
`np.float32`), dtypes (e.g. `np.dtype('float32')`), or objects with a dtype attribute
(e.g. `jnp.float32`, `jnp.int32`).
We may add additional types here in future releases.
JAX Typing Best Practices
-------------------------
When annotating JAX arrays in public API functions, we recommend using :class:`~jax.typing.ArrayLike`
for array inputs, and :class:`~jax.Array` for array outputs.
For example, your function might look like this::
import numpy as np
import jax.numpy as jnp
from jax import Array
from jax.typing import ArrayLike
def my_function(x: ArrayLike) -> Array:
# Runtime type validation, Python 3.10 or newer:
if not isinstance(x, ArrayLike):
raise TypeError(f"Expected arraylike input; got {x}")
# Runtime type validation, any Python version:
if not (isinstance(x, (np.ndarray, Array)) or np.isscalar(x)):
raise TypeError(f"Expected arraylike input; got {x}")
# Convert input to jax.Array:
x_arr = jnp.asarray(x)
# ... do some computation; JAX functions will return Array types:
result = x_arr.sum(0) / x_arr.shape[0]
# return an Array
return result
Most of JAX's public APIs follow this pattern. Note in particular that we recommend JAX functions
to not accept sequences such as :class:`list` or :class:`tuple` in place of arrays, as this can
cause extra overhead in JAX transforms like :func:`~jax.jit` and can behave in unexpected ways with
batch-wise transforms like :func:`~jax.vmap` or :func:`jax.pmap`. For more information on this,
see `Non-array inputs NumPy vs JAX`_
.. _Non-array inputs NumPy vs JAX: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#non-array-inputs-numpy-vs-jax
"""
from jax._src.typing import (
ArrayLike as ArrayLike,
DTypeLike as DTypeLike,
)