Commits

Branch selector

 main

User selector

apaszke

Datepicker

All time

Commit History

Commits on May 4, 2023

• Verify that slices are trivial before discarding them in state primitives

  Show description for 9c5e3f7
  

  apaszke

  authored and

  jax authors

  committed
  9c5e3f7

Commits on Mar 15, 2023

• Optimize canonicalize_shape

  Show description for 1301968
  

  apaszke

  authored and

  jax authors

  committed
  1301968

Commits on Feb 22, 2023

• Slightly increase the tolerance in sparse tests to avoid flakiness

  Show description for 1638313
  

  apaszke

  authored and

  jax authors

  committed
  1638313

Commits on Jan 5, 2023

• Skip gather and reduce scatter grad tests on GPU

  Show description for 6655f2b
  

  apaszke

  authored and

  jax authors

  committed
  6655f2b

• Internal change

  Show description for 904cd4b
  

  apaszke

  authored and

  jax authors

  committed
  904cd4b

• Skip flaky test on TPU

  Show description for f1635ca
  

  apaszke

  authored and

  jax authors

  committed
  f1635ca

Commits on Dec 2, 2022

• Optimize core.find_top_trace

  Show description for bbf22db
  

  apaszke

  authored and

  jax authors

  committed
  bbf22db

Commits on Nov 24, 2022

• Make eager pmap take advantage of pmap cache

  Show description for a711166
  

  apaszke

  authored and

  jax authors

  committed
  a711166

Commits on Nov 23, 2022

• Make testManyArgs actually test pmap with many args

  Show description for 7dd9589
  

  apaszke

  authored and

  jax authors

  committed
  7dd9589

• Shard fft tests to avoid timeouts

  Show description for fe56a19
  

  apaszke

  authored and

  jax authors

  committed
  fe56a19

Commits on Nov 15, 2022

• Transpose all_gather to reduce_scatter

  Show description for d742e6a
  

  apaszke

  authored and

  jax authors

  committed
  d742e6a

Commits on Nov 3, 2022

• Fix the vmap rule for remat_p

  Show description for b0621e3
  

  apaszke

  authored and

  jax authors

  committed
  b0621e3

Commits on Nov 2, 2022

• Don't assume that vmap doesn't introduce constants

  Show description for 94ba43b
  

  apaszke

  authored and

  jax authors

  committed
  94ba43b

Commits on Oct 27, 2022

• Improve undefined axis checks

  Show description for 0fce5be
  apaszke

  committed
  0fce5be

Commits on Oct 26, 2022

• Remove a TODO from the xmap tutorial

  Show description for 6e43ce3
  apaszke

  committed
  6e43ce3

Commits on Oct 14, 2022

• Add support for MANUAL lowering of ppermute

  Show description for 746dd5a
  

  apaszke

  authored and

  jax authors

  committed
  746dd5a

Commits on Oct 11, 2022

• Support MANUAL collectives in top-level xmaps

  Show description for 9c99498
  

  apaszke

  authored and

  jax authors

  committed
  9c99498

Commits on Aug 30, 2022

• Add a correct jaxpr substitution for the cond primitive

  Show description for cbcc77b
  apaszke

  committed
  cbcc77b

Commits on Aug 17, 2022

• Bump xla_client version to gate use_global_device_ids usage.

  Show description for be823b6
  

  apaszke

  authored and

  jax authors

  committed
  be823b6

• Update WORKSPACE

  Show description for 6ed2d22
  apaszke

  committed
  6ed2d22

Commits on Aug 16, 2022

• Fix XLA fallback to avoid checking the mesh conditions

  Show description for 2aea078
  

  apaszke

  authored and

  jax authors

  committed
  2aea078

• Allow collectives in manually sharded computations

  Show description for ffd34d5
  

  apaszke

  authored and

  jax authors

  committed
  ffd34d5

Commits on Aug 5, 2022

• Remove experimental warning from xmap

  Show description for bfb0814
  apaszke

  committed
  bfb0814

Commits on Jul 20, 2022

• Internal change

  Show description for 117da44
  

  apaszke

  authored and

  jax authors

  committed
  117da44

Commits on Jul 5, 2022

• Add support for post_process of xmap in BatchTrace

  Show description for 5777c1e
  

  apaszke

  authored and

  jax authors

  committed
  5777c1e

• Properly count sublevels when tracing xmap body

  Show description for 7439e1b
  apaszke

  committed
  7439e1b

Commits on Feb 22, 2022

• Always treat all mesh axes controlled by xmap as MANUAL

  Show description for 2641f06
  

  apaszke

  authored and

  jax authors

  committed
  2641f06

Commits on Feb 17, 2022

• Fix uninitialized axis_env error when MLIR lowering is disabled

  Show description for 57f4232
  

  apaszke

  authored and

  jax authors

  committed
  57f4232

Commits on Feb 15, 2022

• Try adding support for nesting with_sharding_constraint in MANUAL xmaps

  Show description for a14ccb9
  

  apaszke

  authored and

  jax authors

  committed
  a14ccb9

• Limit the set of unspecified dims to those that are not explicitly converted to MANUAL.

  Show description for b75b0c0
  

  apaszke

  authored and

  jax authors

  committed
  b75b0c0

• Add a partial_eval_jaxpr_custom_rule for xmap

  Show description for c551bed
  

  apaszke

  authored and

  jax authors

  committed
  c551bed

Commits on Feb 7, 2022

• Use aval_out to construct a sharding spec in shard to full

  Show description for 296832e
  

  apaszke

  authored and

  jax authors

  committed
  296832e

• Allow nesting MANUAL-style xmaps in pjits

  Show description for 42cd7ed
  

  apaszke

  authored and

  jax authors

  committed
  42cd7ed

Commits on Feb 4, 2022

• Add experimental support for SPMD lowering of xmap via MANUAL sharding annotations

  Show description for 086a607
  

  apaszke

  authored and

  jax authors

  committed
  086a607

Commits on Dec 10, 2021

• Clarify a comment in ParsedPartitionSpec.eq_given_rank

  Show description for 1b9a9ff
  

  apaszke

  authored and

  jax authors

  committed
  1b9a9ff

Pagination

Previous

Next