Commits

Branch selector

 main

User selector

pschuh

All time

Commit History

Commits on Aug 12, 2024

• Support donating arrays with non-default layouts by setting up XLA donation

  Show description for 734ebd5
  

  pschuh

  authored and

  jax authors

  committed
  734ebd5

  Copy full SHA for 734ebd5

Commits on Aug 9, 2024

• Fix array_test.py when jax_pmap_no_rank_reduction is flipped to true.

  Show description for 4863a56
  

  pschuh

  authored and

  jax authors

  committed
  4863a56

  Copy full SHA for 4863a56

Commits on Jul 31, 2024

• If the product of manual axes is of size 1, then skip emitting

  Show description for d7d9724
  

  pschuh

  authored and

  jax authors

  committed
  d7d9724

  Copy full SHA for d7d9724

Commits on Jul 24, 2024

• Update shard_map(jit) to properly set manual_axes on in_shardings and out_shardings of the nested jit. This avoids a problem where the jit returns {manaual} and then this gets passed to ShardToFull…

  Show description for be6b77c
  

  pschuh

  authored and

  jax authors

  committed
  be6b77c

  Copy full SHA for be6b77c

Commits on Jun 11, 2024

• Allow collectives when only some mesh axes are fully partitioned manually.

  Show description for 4ff7c0f
  

  pschuh

  authored and

  jax authors

  committed
  4ff7c0f

  Copy full SHA for 4ff7c0f

Commits on Jun 5, 2024

• PallasOpsInterpretTest.test_debug_print still flaky, add effects

  Show description for 20c2a45
  

  pschuh

  authored and

  jax authors

  committed
  20c2a45

  Copy full SHA for 20c2a45

• Skip other memories tests that fail on older libtpus.

  Show description for efeb25b
  

  pschuh

  authored and

  jax authors

  committed
  efeb25b

  Copy full SHA for efeb25b

• Bump xla_version and skip crashing test guarded on xla_version

  Show description for d9c5721
  

  pschuh

  authored and

  jax authors

  committed
  d9c5721

  Copy full SHA for d9c5721

Commits on Jun 4, 2024

• Allow disabling compilation cache for particular runtime_types.

  Show description for eba0564
  

  pschuh

  authored and

  jax authors

  committed
  eba0564

  Copy full SHA for eba0564

Commits on May 31, 2024

• Hypothesis isn't compatible with subTest and jax tests sometimes report these

  Show description for 5eb679e
  

  pschuh

  authored and

  jax authors

  committed
  5eb679e

  Copy full SHA for 5eb679e

Commits on May 28, 2024

• Move throwing the error for invalid compute_on into the try-finally context.

  Show description for b47157d
  

  pschuh

  authored and

  jax authors

  committed
  b47157d

  Copy full SHA for b47157d

Commits on May 23, 2024

• Drop refs to australis in build rules and exporter.

  Show description for c1f5a32
  

  pschuh

  authored and

  jax authors

  committed
  c1f5a32

  Copy full SHA for c1f5a32

Commits on May 9, 2024

• Allow nested shard_map.

  Show description for e8ac011
  

  pschuh

  authored and

  jax authors

  committed
  e8ac011

  Copy full SHA for e8ac011

Commits on May 8, 2024

• Cleanup second registration of custom_partitioning callbacks now that

  Show description for e652d62
  

  pschuh

  authored and

  jax authors

  committed
  e652d62

  Copy full SHA for e652d62

Commits on Apr 22, 2024

• Support auto in shard_map.

  Show description for 7ba811e
  

  pschuh

  authored and

  jax authors

  committed
  7ba811e

  Copy full SHA for 7ba811e

Commits on Mar 27, 2024

• Guard host transfers inside pure_callbacks from deadlocking the TPU.

  Show description for 0b09762
  

  pschuh

  authored and

  jax authors

  committed
  0b09762

  Copy full SHA for 0b09762

Commits on Mar 14, 2024

• Propagate effects errors to the results (only if effects are enabled).

  Show description for 9a00721
  

  pschuh

  authored and

  jax authors

  committed
  9a00721

  Copy full SHA for 9a00721

Commits on Feb 13, 2024

• Implement __deepcopy__ for _DimExpr.

  Show description for a93587a
  

  pschuh

  authored and

  jax authors

  committed
  a93587a

  Copy full SHA for a93587a

Commits on Jan 19, 2024

• Return mlir modules instead of XlaComputation from custom_partitioning.

  Show description for 899765e
  

  pschuh

  authored and

  jax authors

  committed
  899765e

  Copy full SHA for 899765e

• [shard_map docs]: Fix some typos

  Show description for 5cf129d
  

  pschuh

  authored and

  jax authors

  committed
  5cf129d

  Copy full SHA for 5cf129d

Commits on Jan 4, 2024

• Add the githash that the jaxlib was built at to __init__.py. This is to allow identifying the githash of nightlies.

  Show description for 23b9c2a
  

  pschuh

  authored and

  jax authors

  committed
  23b9c2a

  Copy full SHA for 23b9c2a

Commits on Dec 18, 2023

• For custom_partitioning, directly emit call when inside of a shard_map.

  Show description for 7ba8622
  

  pschuh

  authored and

  jax authors

  committed
  7ba8622

  Copy full SHA for 7ba8622

Commits on Dec 8, 2023

• Add annotation for donation. This configures XLA's AddBufferDonor directly via donated_args instead of first configuring input_output_aliases. This is best effort anyways, but it works on TPU.

  Show description for ffb115b
  

  pschuh

  authored and

  jax authors

  committed
  ffb115b

  Copy full SHA for ffb115b

Commits on Nov 2, 2023

• Remove temporary flag for forcing arg tuplization of lowered functions.

  Show description for c8b7c1b
  

  pschuh

  authored and

  jax authors

  committed
  c8b7c1b

  Copy full SHA for c8b7c1b

Commits on Oct 27, 2023

• Do not init backends from topology construction, instead directly init the

  Show description for 19c6535
  

  pschuh

  authored and

  jax authors

  committed
  19c6535

  Copy full SHA for 19c6535

Commits on Sep 14, 2023

• Add support for float flags to compiler_options.

  Show description for 2138941
  

  pschuh

  authored and

  jax authors

  committed
  2138941

  Copy full SHA for 2138941

Commits on Sep 7, 2023

• Propagate ad.Zeros to the scan body function for jax.lax.scan for the output 'ys'.

  Show description for bda9292
  

  pschuh

  authored and

  jax authors

  committed
  bda9292

  Copy full SHA for bda9292

Commits on Aug 23, 2023

• Add _manual_axes support to NamedSharding. This is needed because

  Show description for e58ddb7
  

  pschuh

  authored and

  jax authors

  committed
  e58ddb7

  Copy full SHA for e58ddb7

Commits on Aug 16, 2023

• Reverts 75c3457264f9cc117ff09551ce3174d72689fa3d

  Show description for c6f40e2
  

  pschuh

  authored and

  jax authors

  committed
  c6f40e2

  Copy full SHA for c6f40e2

Commits on Aug 11, 2023

• Pad generated sharding specs with None up to ndims to simplify comparing dims

  Show description for 03575c4
  

  pschuh

  authored and

  jax authors

  committed
  03575c4

  Copy full SHA for 03575c4

Commits on Aug 10, 2023

• Make mesh available to custom_partitioning lowering rules.

  Show description for 74bcd65
  

  pschuh

  authored and

  jax authors

  committed
  74bcd65

  Copy full SHA for 74bcd65

Commits on Aug 1, 2023

• Add internal jaxlib function for fetching the topology from

  Show description for 614bbcc
  

  pschuh

  authored and

  jax authors

  committed
  614bbcc

  Copy full SHA for 614bbcc

Commits on Jul 21, 2023

• Add support for int fields to compiler_options.

  Show description for dcad04d
  

  pschuh

  authored and

  jax authors

  committed
  dcad04d

  Copy full SHA for dcad04d

Commits on Jul 17, 2023

• Change visibility for jax_export to public.

  Show description for e9873c4
  

  pschuh

  authored and

  jax authors

  committed
  e9873c4

  Copy full SHA for e9873c4

Commits on Jun 27, 2023

• `jax.lax.collapse` now takes Nones for stop_dimension.

  Show description for 819f731
  

  pschuh

  authored and

  jax authors

  committed
  819f731

  Copy full SHA for 819f731