Commits

Branch selector

 main

User selector

yashk2810

Datepicker

All time

Commit History

Commits on Sep 10, 2024

• Reverts 5e4250e

  Show description for 0b04dd0
  

  yashk2810

  authored and

  jax authors

  committed
  0b04dd0

  Copy full SHA for 0b04dd0

• Prepare for jax 0.4.32 release

  Show description for 5e4250e
  

  yashk2810

  authored and

  jax authors

  committed
  5e4250e

  Copy full SHA for 5e4250e

• Check for jax.Sharding's number of devices instead of py_array.num_shards which looks at IFRT sharding's num_devices to check against global_devices and deciding whether to fall back to pytho…

  Show description for 90892f5
  

  yashk2810

  authored and

  jax authors

  committed
  90892f5

  Copy full SHA for 90892f5

Commits on Sep 9, 2024

• Create optimal order for v5e:8 devices which is [0, 1, 2, 3, 7, 6, 5, 4]

  Show description for d6c3625
  

  yashk2810

  authored and

  jax authors

  committed
  d6c3625

  Copy full SHA for d6c3625

Commits on Sep 5, 2024

• Add compute_on_context_manager to thread local jit state. This is to avoid getting false cache hits

  Show description for a144eb2
  

  yashk2810

  authored and

  jax authors

  committed
  a144eb2

  Copy full SHA for a144eb2

• Don't add a sharding property to ShapedArray if sharding_in_types flag is not switched on.

  Show description for 4c8bed9
  

  yashk2810

  authored and

  jax authors

  committed
  4c8bed9

  Copy full SHA for 4c8bed9

• Fix some tests to use iota_order when creating a mesh

  Show description for 2dd13ce
  

  yashk2810

  authored and

  jax authors

  committed
  2dd13ce

  Copy full SHA for 2dd13ce

Commits on Sep 4, 2024

• Split physical axes by default when device kind is TPU v5 lite to allow for mesh shapes (2, 2) when there are 8 v5e devices on a 4x2 topology.

  Show description for bf66e81
  

  yashk2810

  authored and

  jax authors

  committed
  bf66e81

  Copy full SHA for bf66e81

Commits on Sep 3, 2024

• Rename jtu.create_global_mesh to jtu.create_mesh and use jax.make_mesh inside jtu.create_mesh to get maximum test coverage of the new API.

  Show description for e1b4970
  

  yashk2810

  authored and

  jax authors

  committed
  e1b4970

  Copy full SHA for e1b4970

• Create jax.make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str], devices: Sequence[jax.Device] | None = None) API to make it easier to create a mesh and reduce a ton of boilerplate.

  Show description for 252caeb
  

  yashk2810

  authored and

  jax authors

  committed
  252caeb

  Copy full SHA for 252caeb

Commits on Sep 2, 2024

• Add use_shardy_partitioner to thread local jit state

  Show description for f1e0741
  

  yashk2810

  authored and

  jax authors

  committed
  f1e0741

  Copy full SHA for f1e0741

Commits on Aug 31, 2024

• Reverts changelist 668370165

  Show description for 969dd89
  

  yashk2810

  authored and

  jax authors

  committed
  969dd89

  Copy full SHA for 969dd89

Commits on Aug 30, 2024

• Make input_layouts and output_layouts properties of Compiled. This is equivalent to input_shardings and output_shardings which are also properties on Compiled.

  Show description for 5ba8d6b
  

  yashk2810

  authored and

  jax authors

  committed
  5ba8d6b

  Copy full SHA for 5ba8d6b

• Remove dtype and weak_type from __slots__ of ShapedArray since it comes from UnShapedArray

  Show description for e3110c1
  

  yashk2810

  authored and

  jax authors

  committed
  e3110c1

  Copy full SHA for e3110c1

• Fix failing tests in CI

  Show description for 164b884
  

  yashk2810

  authored and

  jax authors

  committed
  164b884

  Copy full SHA for 164b884

Commits on Aug 29, 2024

• Initial integration of sharding in types in JAX. Currently we just support nary ops in forward only sharding propagation. Currently this functionality is experimental and hidden behind `jax_shard…

  Show description for bcfe95e
  

  2 peopleauthored and

  jax authors

  committed
  bcfe95e

  Copy full SHA for bcfe95e

• Reverts 82c9da0

  Show description for b615266
  

  yashk2810

  authored and

  jax authors

  committed
  b615266

  Copy full SHA for b615266

• Add weak_type to ShapeDtypeStruct because jax.Array also has it and SDS is a duck of jax.Array

  Show description for dd6f0e2
  

  yashk2810

  authored and

  jax authors

  committed
  dd6f0e2

  Copy full SHA for dd6f0e2

• Skip test_layout_donation_mismatching_in_and_out_fails on CPU and GPU

  Show description for 57a9b18
  

  yashk2810

  authored and

  jax authors

  committed
  57a9b18

  Copy full SHA for 57a9b18

Commits on Aug 28, 2024

• Standardize default layout to None in internals (dispatch, lowering and compilation) and non-default layouts to concrete layouts.

  Show description for ef33cf5
  

  yashk2810

  authored and

  jax authors

  committed
  ef33cf5

  Copy full SHA for ef33cf5

• Don't share the same global jit cpp cache between jit and pjit

  Show description for 4695705
  

  yashk2810

  authored and

  jax authors

  committed
  4695705

  Copy full SHA for 4695705

Commits on Aug 27, 2024

• Improve the error message to specify shapes too

  Show description for afff0e0
  

  yashk2810

  authored and

  jax authors

  committed
  afff0e0

  Copy full SHA for afff0e0

Commits on Aug 21, 2024

• Reverts abd442b

  Show description for 08fc5c0
  

  yashk2810

  authored and

  jax authors

  committed
  08fc5c0

  Copy full SHA for 08fc5c0

• Reverts 1e3c079

  Show description for abd442b
  

  yashk2810

  authored and

  jax authors

  committed
  abd442b

  Copy full SHA for abd442b

• Strip primitive params from location info because the amount of metadata included leads to huge HLO size increase and causes compilation cache misses in some other setting too.

  Show description for 1e3c079
  

  yashk2810

  authored and

  jax authors

  committed
  1e3c079

  Copy full SHA for 1e3c079

Commits on Aug 20, 2024

• Generalize global jit cpp cache keys so we can add more keys than the current donate_argnums.

  Show description for 82c9da0
  

  yashk2810

  authored and

  jax authors

  committed
  82c9da0

  Copy full SHA for 82c9da0

• Skip test_in_layouts_jit_jnp_input if xla_extension_version < 282

  Show description for 7cd10d8
  

  yashk2810

  authored and

  jax authors

  committed
  7cd10d8

  Copy full SHA for 7cd10d8

• Skip the global jit cpp cache if in/out_layouts are not None

  Show description for 1ab6279
  

  yashk2810

  authored and

  jax authors

  committed
  1ab6279

  Copy full SHA for 1ab6279

Commits on Aug 19, 2024

• If input layouts are specified via in_shardings to jit and the array that the jitted function is called with is uncommitted, reshard the input array to the layout specified by the user.

  Show description for 6e1c236
  

  yashk2810

  authored and

  jax authors

  committed
  6e1c236

  Copy full SHA for 6e1c236

Commits on Aug 17, 2024

• Match the argument name with the name in Args section in docstring

  Show description for 0c543ae
  

  yashk2810

  authored and

  jax authors

  committed
  0c543ae

  Copy full SHA for 0c543ae

Commits on Aug 15, 2024

• If AbstractMesh is an input to shard_map, then in eager mode require atleast one input to be a NamedSharding not all inputs.

  Show description for 6913551
  

  yashk2810

  authored and

  jax authors

  committed
  6913551

  Copy full SHA for 6913551

Commits on Aug 14, 2024

• Add num_devices to Sharding interface so that it works with NamedSharding containing AbstractMesh too.

  Show description for 229cbae
  

  yashk2810

  authored and

  jax authors

  committed
  229cbae

  Copy full SHA for 229cbae

• Add a devices property to AbstractMesh but raise an error in it. This is to make pytype happy

  Show description for 9a8f0a6
  

  yashk2810

  authored and

  jax authors

  committed
  9a8f0a6

  Copy full SHA for 9a8f0a6

Commits on Aug 13, 2024

• Introduce jax.sharding.AbstractMesh(shape_tuple: tuple[tuple[str, int], ...]) and allow with_sharding_constraint and shard_map to accept an abstract mesh as input (with_sharding_constraint …

  Show description for daa69da
  

  yashk2810

  authored and

  jax authors

  committed
  daa69da

  Copy full SHA for daa69da

• Remove jax_enable_memories conditionals from JAX and remove it from tests too.

  Show description for 4533aea
  

  yashk2810

  authored and

  jax authors

  committed
  4533aea

  Copy full SHA for 4533aea

Pagination

Previous

Next