| HEAD | b06663d | Merge pull request #21358 from hawkinsp:types PiperOrigin-RevId: 636166386 | 22 May 2024, 14:57:25 UTC |
| refs/heads/LenaMartens-patch-1 | f591d0b | Add ensure_compile_time_eval docstring to docs | 14 January 2022, 11:18:40 UTC |
| refs/heads/add-tpu-core-count | 76a7b19 | Set V3 cores to 4 | 15 May 2024, 18:40:57 UTC |
| refs/heads/adx3 | 49b0120 | [adx3] Add string labels for second-order primitive subjaxprs. Subjaxpr labels are helpful for understanding their purpose. | 16 February 2024, 14:55:16 UTC |
| refs/heads/apaszke-debug-gpu | 6b29518 | continue | 27 November 2023, 11:55:58 UTC |
| refs/heads/array_tutorial | aca7e4a | jax.Array tutorial | 16 November 2022, 00:49:17 UTC |
| refs/heads/avals-with-names | fd07c3b | checkpointing the old big bang AVN PR | 08 February 2021, 23:26:31 UTC |
| refs/heads/bazel_tpu | 9d43d6c | tmp | 16 November 2022, 22:45:54 UTC |
| refs/heads/cache_log | c523e3b | Turn down jax_persistent_cache_min_compile_time_secs logging from info to debug. It's very noisy otherwise, since jax usually produces many small computations that aren't cached. | 28 August 2023, 19:49:25 UTC |
| refs/heads/calltf_test | af7d6f0 | Add regression test for hang when jit inputs are pruned. | 17 December 2022, 20:13:21 UTC |
| refs/heads/ci_v3-8 | 8d4b50e | [TPU CI] Run build matrix on v3-8 as well as v4-8 We're seeing failures on v3-8 that don't appear on the current v4-8 testing. v3-8 also exposes 8 devices (vs. v4-8 exposes 4), and some tests needs 8 devices to run. I just added a v3-8 runner VM. Also adds a missing pip install command (I only caught this with a fresh runner since it only needs to be installed once). | 09 December 2022, 22:11:18 UTC |
| refs/heads/cse | ef4a3c6 | polysimp: simplify Mon and Poly classes | 01 April 2019, 04:45:44 UTC |
| refs/heads/debug-nans-test-block-async | e075e64 | block for result in debug_nans_test | 08 October 2020, 20:05:03 UTC |
| refs/heads/dependabot/pip/matplotlib-3.9.0 | f0f44e3 | Bump matplotlib from 3.8.3 to 3.9.0 Bumps [matplotlib](https://github.com/matplotlib/matplotlib) from 3.8.3 to 3.9.0. - [Release notes](https://github.com/matplotlib/matplotlib/releases) - [Commits](https://github.com/matplotlib/matplotlib/compare/v3.8.3...v3.9.0) --- updated-dependencies: - dependency-name: matplotlib dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> | 20 May 2024, 17:41:24 UTC |
| refs/heads/dependabot/pip/pluggy-1.5.0 | 74b05fe | Bump pluggy from 1.4.0 to 1.5.0 Bumps [pluggy](https://github.com/pytest-dev/pluggy) from 1.4.0 to 1.5.0. - [Changelog](https://github.com/pytest-dev/pluggy/blob/main/CHANGELOG.rst) - [Commits](https://github.com/pytest-dev/pluggy/compare/1.4.0...1.5.0) --- updated-dependencies: - dependency-name: pluggy dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> | 20 May 2024, 17:41:32 UTC |
| refs/heads/dependabot/pip/pytest-xdist-3.6.1 | 98cf2d7 | Bump pytest-xdist from 3.5.0 to 3.6.1 Bumps [pytest-xdist](https://github.com/pytest-dev/pytest-xdist) from 3.5.0 to 3.6.1. - [Release notes](https://github.com/pytest-dev/pytest-xdist/releases) - [Changelog](https://github.com/pytest-dev/pytest-xdist/blob/master/CHANGELOG.rst) - [Commits](https://github.com/pytest-dev/pytest-xdist/compare/v3.5.0...v3.6.1) --- updated-dependencies: - dependency-name: pytest-xdist dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> | 20 May 2024, 17:41:37 UTC |
| refs/heads/dependabot/pip/setuptools-70.0.0 | 1128623 | --- updated-dependencies: - dependency-name: setuptools dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> | 22 May 2024, 13:37:15 UTC |
| refs/heads/dependabot/pip/zipp-3.18.2 | ec2770a | Bump zipp from 3.18.1 to 3.18.2 Bumps [zipp](https://github.com/jaraco/zipp) from 3.18.1 to 3.18.2. - [Release notes](https://github.com/jaraco/zipp/releases) - [Changelog](https://github.com/jaraco/zipp/blob/main/NEWS.rst) - [Commits](https://github.com/jaraco/zipp/compare/v3.18.1...v3.18.2) --- updated-dependencies: - dependency-name: zipp dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> | 20 May 2024, 17:41:30 UTC |
| refs/heads/disable_test | 5a54b95 | Disable failing profiler test | 20 July 2023, 22:01:10 UTC |
| refs/heads/dynamic-experiments | 66b182b | wip | 06 December 2019, 15:44:12 UTC |
| refs/heads/dynamic-scoping | e165d18 | wip experiments | 03 July 2019, 16:45:18 UTC |
| refs/heads/effect-types | e74afc1 | rewrites | 18 June 2021, 05:55:33 UTC |
| refs/heads/fix | 008a691 | Remove experimental endpoints and update to point to 0.4.1 | 13 December 2022, 19:34:14 UTC |
| refs/heads/fix-notebooks | 7221e3f | open and re-save notebooks to fix github display | 09 July 2019, 18:33:06 UTC |
| refs/heads/fix-notebooks2 | d18f9a7 | open and resave notebooks in colab | 09 July 2019, 18:47:07 UTC |
| refs/heads/fix_cloud_tpu_check | dfdb74b | Fix test_util.is_cloud_tpu() | 15 November 2023, 00:32:21 UTC |
| refs/heads/fixed-point | a0b92e0 | experiments with a fixed_point implementation | 31 August 2019, 16:07:42 UTC |
| refs/heads/furo | 86648b6 | book | 03 February 2022, 23:40:07 UTC |
| refs/heads/gda1 | f0df7b7 | Fix the display of GDA docs | 03 March 2022, 20:19:04 UTC |
| refs/heads/gda2 | 403eae0 | format nicely | 04 March 2022, 19:59:16 UTC |
| refs/heads/gmap | e69305e | very sketchy vmap collectives Co-authored-by: Adam Paszke <apaszke@google.com> | 30 June 2020, 17:01:45 UTC |
| refs/heads/gnecula-patch-1 | 263a7ff | Update README.md | 22 November 2021, 15:17:09 UTC |
| refs/heads/gpu-determinism-note | 535d984 | add gpu determinism note | 07 November 2020, 23:25:54 UTC |
| refs/heads/hmm-example | 520ba30 | faster normalized version | 03 March 2020, 03:28:48 UTC |
| refs/heads/hoist-consts | 182c462 | hoist constants to parameters in jit-compiled functions Co-authored-by: Matthew Johnson <mattjj@google.com> | 29 May 2020, 00:26:45 UTC |
| refs/heads/ijit | 03bf999 | a group effort at playing with torch inductor (so many coauthors: i wrote the basic ijit and wrote exactly one rule, adam wrote all the rules to make a transformer work plus other improvements, sharad and qiao wrote the model code?) Co-authored-by: Adam Paszke <apaszke@google.com> Co-authored-by: Sharad Vikram <sharadmv@google.com> Co-authored-by: Qiao Zhang <zhangqiaorjc@google.com> | 27 July 2023, 23:24:54 UTC |
| refs/heads/index | a2870a1 | Add jax.Array to the index page | 01 December 2022, 23:38:38 UTC |
| refs/heads/initial-state | a26eeb2 | wip | 22 April 2022, 16:45:22 UTC |
| refs/heads/initial-style-autodidax | c026ec6 | Add a polymorphic list type | 12 February 2024, 14:39:08 UTC |
| refs/heads/issue2263 | 0f115ad | add population_count primitive (needs new jaxlib) fixes #2263 | 19 February 2020, 15:33:26 UTC |
| refs/heads/issue3040 | 8231ec2 | speed up ode_test: make some vjp tests first-order | 20 May 2020, 03:54:23 UTC |
| refs/heads/issue3285 | 389361b | skip fori_loops with zero trip count fixes #3285 | 02 June 2020, 16:20:09 UTC |
| refs/heads/issue3620 | 75c3650 | improve error message for jnp.pad pad_width array fixes #3620 | 01 July 2020, 17:23:34 UTC |
| refs/heads/issue768 | abe4f61 | revive old numpy promotion logic | 28 May 2019, 21:07:29 UTC |
| refs/heads/jax-attrs | cfb9d91 | Trying out JVP. Not feeling eager about eager. Co-authored-by: Matthew Johnson <mattjj@google.com> | 25 January 2024, 21:37:21 UTC |
| refs/heads/jax_release_0.3.6 | 6ba9fb6 | Upgrade the bazel version to 5.1.1 PiperOrigin-RevId: 441338363 | 13 April 2022, 00:47:28 UTC |
| refs/heads/jb/broadcast-tie-in | 7bee52e | add tie_in to broadcast_arrays | 17 November 2019, 02:21:00 UTC |
| refs/heads/jb/debug-strong-zeros | 2f928e7 | wip | 13 October 2019, 00:25:07 UTC |
| refs/heads/jb/partial-eval-scan | 9b296f8 | wip | 10 November 2019, 20:28:37 UTC |
| refs/heads/jb/refs | 686f431 | support refs in api.py transformations | 14 January 2020, 07:39:53 UTC |
| refs/heads/jb/scan-consts | ab282bd | optimize scan with constant-across-time outputs | 05 October 2019, 03:34:33 UTC |
| refs/heads/jb/scan-consts-wip | 18d47d6 | wip, yikes | 05 October 2019, 04:50:45 UTC |
| refs/heads/jb/staxperiments-wip | ee2a92e | pretty happy with this! | 22 October 2019, 01:17:37 UTC |
| refs/heads/jb/tagging | f4f6f16 | fix tag of compile Previously I'd focused on making sure jitted functions called inside tagged ones didn't mess with the tagging tree, by propagating that tree outside argument passing. That doesn't work when we want to tag values inside jitted parts of the tagged function, so this change switches things around to propagate the tagging tree through (e.g. jit) argument passing like ordinary values. | 11 November 2019, 06:49:23 UTC |
| refs/heads/jet | 3e29c9f | add initial jet tests Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> Co-authored-by: David Duvenaud <duvenaud@cs.toronto.edu> Co-authored-by: Dougal Maclaurin <dougalm@google.com> | 02 August 2019, 22:01:45 UTC |
| refs/heads/jet2 | 4f60ed2 | adding div and linear prims (#2338) * adding div and linear prims Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> * fix typo * Add gather rule Co-authored-by: Matthew Johnson <mattjj@csail.mit.edu> Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> Co-authored-by: David Duvenaud <duvenaud@cs.toronto.edu> * jet for reduce_max Co-authored-by: Matthew Johnson <mattjj@csail.mit.edu> Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> Co-authored-by: David Duvenaud <duvenaud@cs.toronto.edu> * linear rule for sub Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> * jet rule for log Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> Co-authored-by: Matthew Johnson <mattjj@csail.mit.edu> Co-authored-by: David Duvenaud <duvenaud@cs.toronto.edu> * add tests Co-authored-by: Jesse Bettencourt <jessebett@cs.toronto.edu> Co-authored-by: Matthew Johnson <mattjj@csail.mit.edu> Co-authored-by: David Duvenaud <duvenaud@cs.toronto.edu> * refactor reduce_max jet rule Co-authored-by: Jacob Kelly <jacob.jin.kelly@gmail.com> Co-authored-by: Matthew Johnson <mattjj@csail.mit.edu> Co-authored-by: David Duvenaud <dduvenaud@seas.harvard.edu> | 02 March 2020, 15:47:52 UTC |
| refs/heads/kokoro-fix | 11fc034 | Fix for spurious deprecation warning in compilation_cache_test | 23 March 2023, 21:48:30 UTC |
| refs/heads/layer-scan-remat | 14d975d | Support call primitives in checkpoint_dots | 10 June 2021, 01:05:30 UTC |
| refs/heads/libtpu_import_fix | 4e34fe0 | Fix libtpu path on older jaxlibs. This is a follow-up to https://github.com/google/jax/commit/b81a3e1fd774ebdbc3015f1bc977bfacb5d4b745. We still need to set TPU_LIBRARY_PATH for jaxlibs that don't support the new mechanism for passing in the libtpu path. | 13 October 2023, 18:20:46 UTC |
| refs/heads/libtpu_install | 1476a85 | Add `--pre` to nightly libtpu pip install command. This is necessary to make sure we pick up the nightly "dev" versions hosted on GCP and not the fake package at https://pypi.org/project/libtpu-nightly/. | 21 February 2023, 23:38:39 UTC |
| refs/heads/linearize | 4773f32 | Bug fixes | 19 May 2020, 15:53:29 UTC |
| refs/heads/lint-action-status | 2ed8762 | post lint/typecheck success/failure commit status | 25 March 2021, 00:14:37 UTC |
| refs/heads/log-resharding | d79f48b | Add logging for pmap input (re)sharding | 13 November 2020, 23:26:10 UTC |
| refs/heads/main | b06663d | Merge pull request #21358 from hawkinsp:types PiperOrigin-RevId: 636166386 | 22 May 2024, 14:57:25 UTC |
| refs/heads/make-custom-vjp-bwd-nones-more-robust | 3e697fb | make the custom_vjp bwd None handling more robust fixes #4673 | 22 October 2020, 22:11:05 UTC |
| refs/heads/masking-revisions | 4a25621 | wip | 09 September 2020, 03:17:05 UTC |
| refs/heads/maxfail | 246614e | Add --maxfail=20 to Cloud TPU CI. This prevents spamming the test output with 100s of failures when something fundamental is broken. Also updates some `python3` commands to use `python` for consistency. | 23 November 2022, 00:47:52 UTC |
| refs/heads/mlir | ae70fb7 | Add generated MLIR _enum_gen.py files to jaxlib wheel. Fixes CI failure after LLVM update. | 25 August 2023, 19:16:15 UTC |
| refs/heads/mutable-array-scan | e63b35d | Add discharge rules for scan with mutable arrays. Move mutable array tests to separate file. Co-authored-by: Matt Johnson <mattjj@google.com> | 29 April 2024, 01:16:13 UTC |
| refs/heads/nanobind | 9a5b808 | Update nanobind version to 1.5.2. | 25 August 2023, 14:45:32 UTC |
| refs/heads/nightly | 74a5c0d | Add nightly to TPU test matrix | 06 April 2023, 23:13:33 UTC |
| refs/heads/no-more-post-process | 2b7abd6 | no more post_process_call / post_process_map | 25 April 2021, 04:41:47 UTC |
| refs/heads/on_demand_ci_jax_gpu | 7307071 | Remove "Publish Test Results" from nightly-ci-multiprocess-gpu.yml | 17 January 2023, 23:31:11 UTC |
| refs/heads/pargmax-primitive | 9fffa86 | add pargmin_p / pargmax_p primitives | 09 February 2021, 04:24:19 UTC |
| refs/heads/pdot | 483d071 | minor refactoring useful for pdot | 25 November 2020, 01:07:30 UTC |
| refs/heads/pjrt_c_api | 473d1c3 | Turn on PJRT C API by default. I forgot that the default setting is actually in jaxlib: https://github.com/openxla/xla/blob/fbe9a80fdb8c429e8a175962459da348cd560a50/xla/python/xla_client.py#L135 To be able to make this change as a jax-only release, I manually set the env var on Cloud TPU if it isn't already set. | 28 March 2023, 20:42:51 UTC |
| refs/heads/pjrt_c_api_tests | ef5e4a4 | Remove 'pjrt_c_api_unimplemented' pytest mark. Instead, we skip tests that the PJRT C API doesn't support. We had this tag for feature development so it was easy to broadly disable, but now we don't expect to need to do that. | 24 March 2023, 20:55:04 UTC |
| refs/heads/pjrt_tpu | 9d1cbc7 | Default to PJRT TPU runtime instead of StreamExecutor on older jaxlibs. I messed up the forwards compat in https://github.com/google/jax/commit/3e50fea29edb6b78426dde511414429f2d2fddf8. The next jaxlib release won't need the env var at all, but jaxlib 0.4.14 and older still do. | 22 August 2023, 20:10:12 UTC |
| refs/heads/pmap | 57c1e4e | progress w/ skye and roy Co-authored-by: Skye Wanderman-Milne <skyewm@google.com> Co-authored-by: Roy Frostig <frostig@google.com> | 31 May 2019, 16:05:51 UTC |
| refs/heads/precision-flag-name-updating | 024cb1b | add more options for jax_default_matmul_precision | 24 March 2021, 23:41:53 UTC |
| refs/heads/pretty-print-improvements | 1500f6f | fix bug | 08 October 2020, 22:29:02 UTC |
| refs/heads/primitives | affd558 | refactor api wrapping code | 03 January 2019, 19:05:40 UTC |
| refs/heads/print-partial-eval-stats | 1e3f700 | print partial eval stats in jit | 22 September 2020, 14:41:56 UTC |
| refs/heads/print_refactor | 365f548 | Merge remote-tracking branch 'origin/print_refactor' into print_refactor | 05 January 2021, 08:51:48 UTC |
| refs/heads/prngkey-linearity-brainstorming | 55e368a | wip | 04 February 2019, 21:33:10 UTC |
| refs/heads/profiler_test | c2d5b46 | profiler_test.py fixes and add coverage to Cloud TPU CI * Add deps to test requirements, including in new `collect-profile-requirements.txt` (to avoid adding tensorflow to `test-requirements.txt`). * Use correct Python executable `ProfilerTest.test_remote_profiler` (`python` sometimes defaults to python2) * Run computations for longer in `ProfilerTest.test_remote_profiler`, othewise `collect_profile` sometimes misses it. | 09 June 2023, 17:58:42 UTC |
| refs/heads/pullrequest | 6c05add | fix things | 15 November 2022, 20:59:00 UTC |
| refs/heads/rafi | 75816ed | For rafi | 17 February 2024, 04:25:12 UTC |
| refs/heads/random-docstring-fix | 64eb46a | Fix RST formatting in random.py docstring | 11 July 2022, 09:51:35 UTC |
| refs/heads/ray | 724804b | easy, just gotta serialize those jaxprs | 08 December 2018, 03:49:13 UTC |
| refs/heads/refs-in-vjps | c33126f | fix | 12 April 2024, 21:25:38 UTC |
| refs/heads/rejames3 | 4a4d7f7 | draft! | 06 July 2021, 19:16:42 UTC |
| refs/heads/restore_jax_remat_opt_barrier_4.2.0 | 03c8954 | Fix `'jax.lax' has no attribute 'empty'` error" | 08 February 2023, 22:51:09 UTC |
| refs/heads/revive-leak-checker | 9b12763 | revive the tracer leak checker The tracer leak checker never worked with the jit (or pmap) compilation cache because 1. it relies on Python reference counting (via a weakref mechanism) to check that there are no more references to a MasterTrace (e.g. from any Tracers associated with it) once a trace is finished, but 2. the compilation caches (i.e. linear_util.cache) can include in their cache key the transforms stacked on the corresponding WrappedFun being jitted, and transforms (specifically trace_to_subjaxpr in partial_eval.py) can include MasterTraces. Hence the cache keys included references to the MasterTraces, defeating the leak checking mechanism. This commit just makes an equal copy of any MasterTraces when building the cache key. MasterTraces are already hashable, with equality defined based on just their level and trace type. MasterTraces are only compared by identity in core.full_raise, and then only to determine if a sublift is needed (because a trace is encountering one of its own tracers but from inside at least one additional level of call scoping). That's not an issue for jit because of the caching, but it could be another issue for calls; TODO figure that out. Co-authored-by: James Bradbury <jekbradbury@google.com> | 23 May 2020, 02:17:44 UTC |
| refs/heads/rvjp | 036b5b5 | Clean up prints an asserts | 20 November 2019, 21:27:39 UTC |
| refs/heads/sc | 88408e1 | Remove stale references to //jaxlib:setup.cfg in Bazel build. Fixes broken jaxlib wheel build. | 03 September 2023, 19:18:25 UTC |
| refs/heads/setupcfg | 88408e1 | Remove stale references to //jaxlib:setup.cfg in Bazel build. Fixes broken jaxlib wheel build. | 03 September 2023, 19:18:25 UTC |
| refs/heads/sharadmv-patch-1 | e1da5c7 | Update CHANGELOG.md | 03 August 2022, 05:28:52 UTC |
| refs/heads/sharadmv-patch-2 | 1a8a8a5 | Fix example in `pjit` docstring | 22 September 2022, 19:55:55 UTC |
| refs/heads/sharding_docstrings | aea7936 | Sharding docstrings | 12 November 2022, 00:33:53 UTC |
