Commits

Branch selector

 main

User selector

mattjj

All time

Commit History

Commits on Aug 19, 2024

• fix basic scan bug with attrs

  mattjj

  committed
  ef82cb2

  Copy full SHA for ef82cb2

Commits on Aug 8, 2024

• fix #22944

  mattjj

  committed
  44ae9b3

  Copy full SHA for 44ae9b3

Commits on Aug 3, 2024

• improve while_loop carry pytree/type mismatch errors

  Show description for bdcd358
  mattjj

  committed
  bdcd358

  Copy full SHA for bdcd358

Commits on Jul 31, 2024

• add temporary flag to suppress an error message, to unblock a user

  mattjj

  committed
  abfb1ce

  Copy full SHA for abfb1ce

Commits on Jul 25, 2024

• deflake cond memory leak regression test

  mattjj

  committed
  f72a3f8

  Copy full SHA for f72a3f8

• remove pdot and xeinsum (since xmap is gone)

  mattjj

  committed
  88d1cd7

  Copy full SHA for 88d1cd7

• remove inlined jax.nn.initializers definitions, resolving TODO of levskaya et al

  Show description for c8ea86c
  

  mattjj

  authored and

  jax authors

  committed
  c8ea86c

  Copy full SHA for c8ea86c

• remove named_shapes (since xmap is now gone)

  mattjj

  committed
  3f9eb40

  Copy full SHA for 3f9eb40

Commits on Jul 23, 2024

• fix memory leak in cond jaxpr tracig

  Show description for 8db862c
  mattjj

  committed
  8db862c

  Copy full SHA for 8db862c

• make grad-of-pmap with out_axes=None raise NotImplementedError

  Show description for 405872d
  mattjj

  committed
  405872d

  Copy full SHA for 405872d

Commits on Jul 22, 2024

• add test for #17691

  Show description for 3beb3d5
  mattjj

  committed
  3beb3d5

  Copy full SHA for 3beb3d5

• [shard_map] fix psum rewrite rule's pbroadcast logic

  Show description for f7cef92
  mattjj

  committed
  f7cef92

  Copy full SHA for f7cef92

Commits on Jul 20, 2024

• [shard_map] shard_map check_rep=True rules for custom_linear_solve

  Show description for 173794b
  mattjj

  committed
  173794b

  Copy full SHA for 173794b

• skip _origin_msg invar debug info if invar_pos/arg_info is malformed

  Show description for c5fd3b0
  mattjj

  committed
  c5fd3b0

  Copy full SHA for c5fd3b0

• make core.as_named_shape treat `int` like `tuple[int]`

  Show description for 83dfed1
  mattjj

  committed
  83dfed1

  Copy full SHA for 83dfed1

Commits on Jul 11, 2024

• make remat reduce precision of saved values to avoid xla excess precision

  Show description for f498d7a
  

  mattjj

  and

  hawkinsp

  committed
  f498d7a

  Copy full SHA for f498d7a

Commits on Jul 2, 2024

• add more info to xla_computation deprecation warning

  mattjj

  committed
  bd166e1

  Copy full SHA for bd166e1

Commits on Jul 1, 2024

• prototyping improving vmap out_axes error

  Show description for 987194d
  

  2 peopleauthored and

  cgarciae

  committed
  987194d

  Copy full SHA for 987194d

Commits on Jun 26, 2024

• tweak dynamic trace state to only depend on level int, not MainTrace

  mattjj

  committed
  275ddad

  Copy full SHA for 275ddad

• [shard-map] allow in_specs=None to work with arbitrary objects

  Show description for df90711
  mattjj

  committed
  df90711

  Copy full SHA for df90711

• fix error message path since #22069 allowed Literals in jaxpr builder eqns

  mattjj

  committed
  ad4e5ec

  Copy full SHA for ad4e5ec

• remove double dots from an error message

  mattjj

  committed
  8564b55

  Copy full SHA for 8564b55

Commits on Jun 24, 2024

• faster inline jaxpr, but maybe we shouldnt inline at all

  mattjj

  committed
  789a0c7

  Copy full SHA for 789a0c7

Commits on Jun 5, 2024

• make cond forward inputs to outputs, reduces vmap lifting

  Show description for 7c12570
  

  mattjj

  and

  Cristian Garcia

  committed
  7c12570

  Copy full SHA for 7c12570

Commits on May 31, 2024

• [custom_vjp] allow bwd rule to produce top-level list (not just tuple)

  mattjj

  committed
  2f4cc6e

  Copy full SHA for 2f4cc6e

Commits on May 30, 2024

• fix error message for vjp arguments

  mattjj

  committed
  10d285d

  Copy full SHA for 10d285d

• add error checks for vmap spmd_axis_name

  mattjj

  committed
  3984d82

  Copy full SHA for 3984d82

Commits on May 29, 2024

• fix weak key cache stuff

  Show description for a4622b6
  

  mattjj

  and

  dougalm

  committed
  a4622b6

  Copy full SHA for a4622b6

Commits on May 25, 2024

• avoid singleton dim in scan lowering when unroll==1

  mattjj

  committed
  a24b738

  Copy full SHA for a24b738

• add pjit forwarding rule

  Show description for 0a693fa
  

  mattjj

  and

  froystig

  committed
  0a693fa

  Copy full SHA for 0a693fa

Commits on May 4, 2024

• [shard_map] better fix for spmd_axis_name issues with shmap residuals

  Show description for 7a87010
  mattjj

  committed
  7a87010

  Copy full SHA for 7a87010

Commits on May 2, 2024

• [omnitracing] partially un-regress dispatch time

  mattjj

  committed
  e4c76c9

  Copy full SHA for e4c76c9

• fix vmap-of-grad-of-shmap axis name reuse bug

  Show description for b7b2645
  mattjj

  committed
  b7b2645

  Copy full SHA for b7b2645

Commits on Apr 30, 2024

• [omnitracing] pop frames, see if anything downstream breaks

  mattjj

  committed
  7af2b3c

  Copy full SHA for 7af2b3c

• skip pallas/gmm_test.py if we don't have hypothesis

  mattjj

  committed
  be499e0

  Copy full SHA for be499e0