Commits

Branch selector

 main

User selector

apaszke

Datepicker

All time

Commit History

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

  Copy full SHA for 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

  Copy full SHA for 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

  Copy full SHA for 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

  Copy full SHA for b75b0c0

• Add a partial_eval_jaxpr_custom_rule for xmap

  Show description for c551bed
  

  apaszke

  authored and

  jax authors

  committed
  c551bed

  Copy full SHA for 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

  Copy full SHA for 296832e

• Allow nesting MANUAL-style xmaps in pjits

  Show description for 42cd7ed
  

  apaszke

  authored and

  jax authors

  committed
  42cd7ed

  Copy full SHA for 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

  Copy full SHA for 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

  Copy full SHA for 1b9a9ff

Commits on Oct 15, 2021

• Enable batcher and batched collective rules for tiled all gathers

  Show description for 49d9aff
  apaszke

  committed
  49d9aff

  Copy full SHA for 49d9aff

Commits on Oct 12, 2021

• Fix while loop batching rule for loops with constant bodies

  Show description for b8dc8ca
  apaszke

  committed
  b8dc8ca

  Copy full SHA for b8dc8ca

Commits on Oct 11, 2021

• Disallow jax2tf translation of multi-process pjits

  Show description for dad29d3
  apaszke

  committed
  dad29d3

  Copy full SHA for dad29d3

Commits on Oct 5, 2021

• Change semantics of positional shapes depending on pjit nesting

  Show description for 1520fa2
  

  apaszke

  authored and

  jax authors

  committed
  1520fa2

  Copy full SHA for 1520fa2

Commits on Oct 4, 2021

• Strengthen resource type checking to verify physical mesh consistency

  Show description for 490d049
  

  apaszke

  authored and

  jax authors

  committed
  490d049

  Copy full SHA for 490d049

Commits on Oct 1, 2021

• Fix xmap error message when checking axis specs

  Show description for 6c36c3e
  apaszke

  committed
  6c36c3e

  Copy full SHA for 6c36c3e

• Keep axis_env initialized during jaxpr_subcomp

  Show description for 08685ef
  apaszke

  committed
  08685ef

  Copy full SHA for 08685ef

Commits on Sep 30, 2021

• Raise a clear error message when pjit is used with CPU devices

  Show description for 004d855
  

  apaszke

  authored and

  jax authors

  committed
  004d855

  Copy full SHA for 004d855

Commits on Sep 29, 2021

• Improve the error message raised when local mesh is not a continuous subcube

  Show description for 2a4648b
  

  apaszke

  authored and

  jax authors

  committed
  2a4648b

  Copy full SHA for 2a4648b

• Add reverse-mode AD support for pjit

  Show description for 22dce0f
  

  apaszke

  authored and

  jax authors

  committed
  22dce0f

  Copy full SHA for 22dce0f

Commits on Sep 27, 2021

• Use shaped_abstractify to allow abstract arguments in .lower methods

  Show description for 484c7e0
  

  apaszke

  authored and

  jax authors

  committed
  484c7e0

  Copy full SHA for 484c7e0

Commits on Sep 20, 2021

• Make sure that xmap raises a clear error when using an undefined resource name

  Show description for 45c1a1b
  

  apaszke

  authored and

  jax authors

  committed
  45c1a1b

  Copy full SHA for 45c1a1b

Commits on Sep 13, 2021

• Expose allow_spmd_sharding_propagation_to_output in Python

  Show description for 3218b06
  

  apaszke

  authored and

  jax authors

  committed
  3218b06

  Copy full SHA for 3218b06

Commits on Sep 10, 2021

• Add new lowering APIs for xmap

  Show description for cc2efbf
  

  apaszke

  authored and

  jax authors

  committed
  cc2efbf

  Copy full SHA for cc2efbf

• Cache used_axis_names calls for jaxprs

  Show description for c845d15
  

  apaszke

  authored and

  jax authors

  committed
  c845d15

  Copy full SHA for c845d15

Commits on Sep 9, 2021

• Add new lowering APIs for jit

  Show description for 00528a4
  

  apaszke

  authored and

  jax authors

  committed
  00528a4

  Copy full SHA for 00528a4

Commits on Sep 8, 2021

• New lowering APIs for pjit

  Show description for 5b4757d
  

  apaszke

  authored and

  jax authors

  committed
  5b4757d

  Copy full SHA for 5b4757d

• Remove axis name from named_shape when unmapping avals

  Show description for 1158530
  

  apaszke

  authored and

  jax authors

  committed
  1158530

  Copy full SHA for 1158530

Commits on Sep 7, 2021

• Don't cast out_axis_resources to a tuple automatically

  Show description for 0c03e98
  

  apaszke

  authored and

  jax authors

  committed
  0c03e98

  Copy full SHA for 0c03e98

• Ensure that named axes consistently refer to global axis sizes in xmap

  Show description for 0636f49
  

  apaszke

  authored and

  jax authors

  committed
  0636f49

  Copy full SHA for 0636f49

Commits on Aug 30, 2021

• Revert of "change while loop batching fixed point condition"

  Show description for e410085
  

  apaszke

  authored and

  jax authors

  committed
  e410085

  Copy full SHA for e410085

Commits on Jul 19, 2021

• Add an option to strictly enforce sharding implies by named axes

  Show description for d25f4b3
  

  apaszke

  authored and

  jax authors

  committed
  d25f4b3

  Copy full SHA for d25f4b3

Commits on Jul 15, 2021

• Make maps.EXPERIMENTAL_SPMD_LOWERING into a jax.config flag

  Show description for e987f6f
  apaszke

  committed
  e987f6f

  Copy full SHA for e987f6f

• Add support for reverse-mode AD of xmap

  apaszke

  committed
  478828a

  Copy full SHA for 478828a

• Add axis and tiled options to lax.all_gather.

  Show description for 64510bd
  

  apaszke

  authored and

  jax authors

  committed
  64510bd

  Copy full SHA for 64510bd

• Make maps.EXPERIMENTAL_SPMD_LOWERING into a jax.config flag

  Show description for 8bc6e7f
  

  apaszke

  authored and

  jax authors

  committed
  8bc6e7f

  Copy full SHA for 8bc6e7f

Pagination

Previous

Next