Commits

Branch selector

 main

User selector

mattjj

Datepicker

All time

Commit History

Commits on Feb 13, 2023

• fix minor broadcasting bug

  Show description for 96c558d
  

  mattjj

  and

  apaszke

  committed
  96c558d

  Copy full SHA for 96c558d

Commits on Feb 10, 2023

• generalize vmap spmd_axis_name to accept tuples of axis names

  Show description for 9538bc3
  mattjj

  committed
  9538bc3

  Copy full SHA for 9538bc3

Commits on Feb 9, 2023

• simpler pretty-print for pjit, tweak custom pp rule signature

  mattjj

  committed
  a964dc3

  Copy full SHA for a964dc3

• [shard-map] add vmap spmd_axis_name support, fix vmap rule bug

  mattjj

  committed
  6fb3ace

  Copy full SHA for 6fb3ace

Commits on Feb 8, 2023

• [shard-map] if check_rep=False, don't call rep rules in eager

  mattjj

  committed
  1a03f34

  Copy full SHA for 1a03f34

• [shard-map] add remat support, very basic test

  mattjj

  committed
  58d3f55

  Copy full SHA for 58d3f55

Commits on Feb 7, 2023

• [shard_map] add a lot of trivial rules

  mattjj

  committed
  198bfe3

  Copy full SHA for 198bfe3

• [shard_map] add rep rule for axis_index, trivial test

  mattjj

  committed
  6db3f48

  Copy full SHA for 6db3f48

Commits on Feb 3, 2023

• shard_map (shmap) prototype and JEP

  Show description for ff1e9b3
  

  3 peoplecommitted
  ff1e9b3

  Copy full SHA for ff1e9b3

• minor tweaks to type annotations, specialize code on those types

  Show description for 644d3b6
  mattjj

  committed
  644d3b6

  Copy full SHA for 644d3b6

Commits on Feb 2, 2023

• skip custom_jvp/vjp tests which dont work with initial-style staging

  Show description for cd615b6
  mattjj

  committed
  cd615b6

  Copy full SHA for cd615b6

Commits on Feb 1, 2023

• checkify: cache jaxpr formation so we don't always retrace

  mattjj

  committed
  684846b

  Copy full SHA for 684846b

Commits on Jan 20, 2023

• update pjit test

  mattjj

  committed
  358775f

  Copy full SHA for 358775f

• specialize tree prefix error message for list/tuple

  mattjj

  committed
  cea2b6b

  Copy full SHA for cea2b6b

Commits on Jan 16, 2023

• cond transpose, use UndefinedPrimal not `linear` for transpose inputs

  mattjj

  committed
  e516d41

  Copy full SHA for e516d41

Commits on Jan 15, 2023

• fix transcription error for initial_step_size (thanks @AlbertMitjans)

  Show description for 1da24b6
  mattjj

  committed
  1da24b6

  Copy full SHA for 1da24b6

Commits on Jan 11, 2023

• add pjit partial_eval_jaxpr_custom rule

  Show description for 8b58530
  

  mattjj

  and

  yashk2810

  committed
  8b58530

  Copy full SHA for 8b58530

Commits on Jan 4, 2023

• tweak remat named policies

  mattjj

  committed
  4f7cf62

  Copy full SHA for 4f7cf62

Commits on Dec 21, 2022

• tweak dot_general pretty-printing rule to suppress default params

  mattjj

  committed
  c2d9b5c

  Copy full SHA for c2d9b5c

• relax tanh test tols for upcoming xla change

  mattjj

  committed
  6ba0ef6

  Copy full SHA for 6ba0ef6

Commits on Dec 20, 2022

• tweak print_saved_residuals

  mattjj

  committed
  580fdb6

  Copy full SHA for 580fdb6

Commits on Dec 11, 2022

• in jax.Array notebook, polish beginning and tweak title and some wording

  mattjj

  committed
  1185c89

  Copy full SHA for 1185c89

Commits on Nov 8, 2022

• tighten up vmap w/ piles: require pile_axis in_axes/out_axes

  mattjj

  committed
  0b463ef

  Copy full SHA for 0b463ef

Commits on Nov 7, 2022

• add a basic prototype of piles, behind jax_dynamic_shapes

  Show description for f2f2faa
  

  3 peoplecommitted
  f2f2faa

  Copy full SHA for f2f2faa

Commits on Nov 5, 2022

• fix jax.random.logits shape argument

  Show description for 190204f
  mattjj

  committed
  190204f

  Copy full SHA for 190204f

Commits on Nov 3, 2022

• improve error when f_vjp gets more than one argument

  Show description for 4033007
  mattjj

  committed
  4033007

  Copy full SHA for 4033007

• fix comparison table in random docs

  Show description for 478bd3e
  

  mattjj

  and

  froystig

  committed
  478bd3e

  Copy full SHA for 478bd3e

Commits on Oct 29, 2022

• integrate new (partitionable, count-space-exhaustive) counts generation

  mattjj

  committed
  213d2c8

  Copy full SHA for 213d2c8

Commits on Oct 28, 2022

• fix ci failure by skipping tests on gpu

  mattjj

  committed
  0203534

  Copy full SHA for 0203534

• make leak checker errors explain why objects are alive

  Show description for 6ebf44a
  

  3 peoplecommitted
  6ebf44a

  Copy full SHA for 6ebf44a

Commits on Oct 26, 2022

• [dynamic-shapes] tweak jnp.canonicalize_shape logic

  Show description for 7e34181
  mattjj

  committed
  7e34181

  Copy full SHA for 7e34181

Commits on Oct 25, 2022

• hopefully fix pjit bug

  Show description for 612bb17
  

  mattjj

  and

  yashk2810

  committed
  612bb17

  Copy full SHA for 612bb17

• tweaks to DevicesSharding

  Show description for 95eb424
  mattjj

  committed
  95eb424

  Copy full SHA for 95eb424

Commits on Oct 21, 2022

• fix

  mattjj

  committed
  8e8ae84

  Copy full SHA for 8e8ae84

• put back some unsafe_maps

  mattjj

  committed
  f76fc01

  Copy full SHA for f76fc01

Pagination

Previous

Next