Commits

Branch selector

 main

User selector

mattjj

Datepicker

All time

Commit History

Commits on Feb 16, 2024

• prototype 'remat' in hlo metadata

  mattjj

  committed
  08dfb11

  Copy full SHA for 08dfb11

• [scan] don't traverse body jaxpr in lowering

  Show description for 5ead7a6
  mattjj

  committed
  5ead7a6

  Copy full SHA for 5ead7a6

Commits on Feb 14, 2024

• [attrs] allow passing a jax-attrs object to jit functions

  Show description for a45cc43
  mattjj

  committed
  a45cc43

  Copy full SHA for a45cc43

Commits on Feb 13, 2024

• [attrs] add a jvp function with attrs support

  Show description for cdb466c
  

  mattjj

  and

  dougalm

  committed
  cdb466c

  Copy full SHA for cdb466c

Commits on Feb 1, 2024

• basic scan support

  Show description for 039e8df
  mattjj

  committed
  039e8df

  Copy full SHA for 039e8df

Commits on Jan 30, 2024

• add getstate/setstate in pjit transpose, for bwd pass effects

  Show description for 6c2d9c7
  

  mattjj

  and

  froystig

  committed
  6c2d9c7

  Copy full SHA for 6c2d9c7

Commits on Jan 28, 2024

• add test

  mattjj

  committed
  22160df

  Copy full SHA for 22160df

• integrate attrs in jax.jit

  Show description for 4a8babb
  

  mattjj

  and

  dougalm

  committed
  4a8babb

  Copy full SHA for 4a8babb

Commits on Jan 27, 2024

• make hypothesis dependence optional

  mattjj

  committed
  54d7d5c

  Copy full SHA for 54d7d5c

Commits on Jan 23, 2024

• [shard-map] fix disable_jit-of-shmap

  mattjj

  committed
  5275ef6

  Copy full SHA for 5275ef6

• [shard-map] fix typos in docs

  mattjj

  committed
  a61f780

  Copy full SHA for a61f780

Commits on Jan 19, 2024

• Allow unhashable callables in jax.eval_shape.

  Show description for 0d4f200
  

  mattjj

  authored and

  jax authors

  committed
  0d4f200

  Copy full SHA for 0d4f200

Commits on Jan 18, 2024

• [shard-map] add user tutorial

  mattjj

  committed
  8b219d5

  Copy full SHA for 8b219d5

Commits on Jan 16, 2024

• [shard-map] add approx_top_k replication rule

  mattjj

  committed
  30c0fc4

  Copy full SHA for 30c0fc4

Commits on Jan 11, 2024

• improve block matrix alignment in pallas docs

  mattjj

  committed
  efe78c5

  Copy full SHA for efe78c5

Commits on Jan 2, 2024

• [shard-map] improve error message when a custom_vjp bwd has extra psum

  mattjj

  committed
  12e57de

  Copy full SHA for 12e57de

Commits on Dec 27, 2023

• add jax.explain_cache_misses tracing cache miss explanations

  Show description for 9112dce
  

  mattjj

  and

  yashk2810

  committed
  9112dce

  Copy full SHA for 9112dce

Commits on Dec 24, 2023

• add test illustrating scale dtypes in scan

  mattjj

  committed
  0ea9a71

  Copy full SHA for 0ea9a71

• revive tangent dtypes test broken by d5646db

  Show description for d80d7a7
  mattjj

  committed
  d80d7a7

  Copy full SHA for d80d7a7

Commits on Dec 22, 2023

• partial rollback of #19096 due to internal breakage (relying on jax internals)

  Show description for d5646db
  

  mattjj

  authored and

  jax authors

  committed
  d5646db

  Copy full SHA for d5646db

• tweaks to enable adding custom tangent dtypes

  Show description for 05da18a
  mattjj

  committed
  05da18a

  Copy full SHA for 05da18a

• del add_any_p and zeros_like_p, replace aval-dispatched traceable

  mattjj

  committed
  be3ca50

  Copy full SHA for be3ca50

Commits on Dec 21, 2023

• in partial_eval_custom rule for pjit, cache ClosedJaxpr creation

  Show description for e31018b
  

  mattjj

  and

  yashk2810

  committed
  e31018b

  Copy full SHA for e31018b

• remove use of cast in ad_util

  Show description for 6863569
  mattjj

  committed
  6863569

  Copy full SHA for 6863569

Commits on Dec 20, 2023

• revise logic for tangent types of extended dtypes

  Show description for ec7d28c
  mattjj

  committed
  ec7d28c

  Copy full SHA for ec7d28c

• make hypothesis use in all_gather_test.py optional

  Show description for d7cceb1
  mattjj

  committed
  d7cceb1

  Copy full SHA for d7cceb1

• handle convert_element_type(complex -> real) in constant folding

  Show description for 325a008
  mattjj

  committed
  325a008

  Copy full SHA for 325a008

Commits on Dec 15, 2023

• fix a bug with eager pmap + vmap + custom_jvp interaction

  Show description for b70ac90
  mattjj

  committed
  b70ac90

  Copy full SHA for b70ac90

Commits on Dec 13, 2023

• [shard-map] register cumsum et al with generic rules

  Show description for 4ba6bd5
  mattjj

  committed
  4ba6bd5

  Copy full SHA for 4ba6bd5

• [shard-map] fix varargs error message bug

  Show description for 2bff2f4
  

  mattjj

  and

  chaserileyroberts

  committed
  2bff2f4

  Copy full SHA for 2bff2f4

Commits on Dec 9, 2023

• don't call lax.xeinsum from jnp.einsum when str contains '{'

  Show description for 25eb913
  mattjj

  committed
  25eb913

  Copy full SHA for 25eb913

• remove _use_xeinsum from jnp.einsum api

  Show description for 9a1a09c
  mattjj

  committed
  9a1a09c

  Copy full SHA for 9a1a09c

• fix gc bug in PjitFunction (not clearing/visiting all children)

  Show description for 7010650
  

  3 peopleauthored and

  jax authors

  committed
  7010650

  Copy full SHA for 7010650

Commits on Dec 8, 2023

• must flatten defaultdict in key-sorted order, like regular dicts

  mattjj

  committed
  d2fcf27

  Copy full SHA for d2fcf27

Commits on Dec 7, 2023

• improve an error message during Mesh creation

  mattjj

  committed
  64cb53f

  Copy full SHA for 64cb53f

Pagination

Previous

Next