Commits

Branch selector

 main

User selector

mattjj

Datepicker

All time

Commit History

Commits on Sep 7, 2024

• tweak readme title to be more about what jax can do for you, dear user

  mattjj

  committedSep 7, 2024
  8ab5031

Commits on Sep 2, 2024

• tweak shmap implementation to work better with leak checker

  mattjj

  committedSep 2, 2024
  f2bef6b

Commits on Aug 23, 2024

• add experimental jax.no_tracing context manager

  mattjj

  committedAug 23, 2024
  670a648

Commits on Aug 19, 2024

• fix basic scan bug with attrs

  mattjj

  committedAug 19, 2024
  ef82cb2

Commits on Aug 8, 2024

• fix google#22944

  mattjj

  committedAug 8, 2024
  44ae9b3

Commits on Aug 3, 2024

• improve while_loop carry pytree/type mismatch errors

  mattjj

  committedAug 3, 2024
  bdcd358

Commits on Jul 31, 2024

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

  mattjj

  committedJul 31, 2024
  abfb1ce

Commits on Jul 25, 2024

• deflake cond memory leak regression test

  mattjj

  committedJul 25, 2024
  f72a3f8

• remove pdot and xeinsum (since xmap is gone)

  mattjj

  committedJul 25, 2024
  88d1cd7

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

  

  mattjj

  authored and

  jax authors

  committedJul 25, 2024
  c8ea86c

• remove named_shapes (since xmap is now gone)

  mattjj

  committedJul 25, 2024
  3f9eb40

Commits on Jul 23, 2024

• fix memory leak in cond jaxpr tracig

  mattjj

  committedJul 23, 2024
  8db862c

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

  mattjj

  committedJul 23, 2024
  405872d

Commits on Jul 22, 2024

• add test for google#17691

  mattjj

  committedJul 22, 2024
  3beb3d5

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

  mattjj

  committedJul 22, 2024
  f7cef92

Commits on Jul 20, 2024

• [shard_map] shard_map check_rep=True rules for custom_linear_solve

  mattjj

  committedJul 20, 2024
  173794b

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

  mattjj

  committedJul 20, 2024
  c5fd3b0

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

  mattjj

  committedJul 20, 2024
  83dfed1

Commits on Jul 11, 2024

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

  

  mattjj

  and

  hawkinsp

  committedJul 11, 2024
  f498d7a

Commits on Jul 2, 2024

• add more info to xla_computation deprecation warning

  mattjj

  committedJul 2, 2024
  bd166e1

Commits on Jul 1, 2024

• prototyping improving vmap out_axes error

  

  2 peopleauthored and

  cgarciae

  committedJul 1, 2024
  987194d

Commits on Jun 26, 2024

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

  mattjj

  committedJun 26, 2024
  275ddad

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

  mattjj

  committedJun 26, 2024
  df90711

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

  mattjj

  committedJun 26, 2024
  ad4e5ec

• remove double dots from an error message

  mattjj

  committedJun 26, 2024
  8564b55

Commits on Jun 24, 2024

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

  mattjj

  committedJun 24, 2024
  789a0c7

Commits on Jun 5, 2024

• make cond forward inputs to outputs, reduces vmap lifting

  

  mattjj

  and

  Cristian Garcia

  committedJun 5, 2024
  7c12570

Commits on May 31, 2024

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

  mattjj

  committedMay 31, 2024
  2f4cc6e

Commits on May 30, 2024

• fix error message for vjp arguments

  mattjj

  committedMay 30, 2024
  10d285d

• add error checks for vmap spmd_axis_name

  mattjj

  committedMay 30, 2024
  3984d82

Commits on May 29, 2024

• fix weak key cache stuff

  

  mattjj

  and

  dougalm

  committedMay 29, 2024
  a4622b6

Commits on May 25, 2024

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

  mattjj

  committedMay 25, 2024
  a24b738

• add pjit forwarding rule

  

  mattjj

  and

  froystig

  committedMay 25, 2024
  0a693fa

Commits on May 4, 2024

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

  mattjj

  committedMay 4, 2024
  7a87010

Commits on May 2, 2024

• [omnitracing] partially un-regress dispatch time

  mattjj

  committedMay 2, 2024
  e4c76c9

Pagination

PreviousNext