Commits

Branch selector

 main

User selector

yashk2810

Datepicker

All time

Commit History

Commits on Apr 1, 2022

• Give auto sharder the mesh information specifically the mesh_shape and the devices

  Show description for 0eaeff6
  

  yashk2810

  authored and

  jax authors

  committed
  0eaeff6

  Copy full SHA for 0eaeff6

• Give auto sharder the mesh information specifically the mesh_shape and the devices

  Show description for 7b7458b
  

  yashk2810

  authored and

  jax authors

  committed
  7b7458b

  Copy full SHA for 7b7458b

• Fix the breakage by including --experimental_cc_shared_library as done by TF.

  Show description for aa5d6b4
  

  yashk2810

  authored and

  jax authors

  committed
  aa5d6b4

  Copy full SHA for aa5d6b4

• First attempt to enable auto-sharding. This CL adds support for GDA (no SDA support yet).

  Show description for 8ca8f74
  

  yashk2810

  authored and

  jax authors

  committed
  8ca8f74

  Copy full SHA for 8ca8f74

Commits on Mar 31, 2022

• If jax_parallel_functions_output_gda flag is set to True, then all outputs are GDA. In abstract_eval rule of pjit, don't convert global avals to local avals if the `jax_parallel_functions_outpu…

  Show description for 0bcd447
  

  yashk2810

  authored and

  jax authors

  committed
  0bcd447

  Copy full SHA for 0bcd447

Commits on Mar 29, 2022

• Speed up GDA initialization by making local_shards lazy and hiding checks behind config.jax_enable_checks flag.

  Show description for e08bc27
  

  yashk2810

  authored and

  jax authors

  committed
  e08bc27

  Copy full SHA for e08bc27

• Fix the mesh_axes of benchmarks to be Pspecs

  Show description for a68b0f3
  

  yashk2810

  authored and

  jax authors

  committed
  a68b0f3

  Copy full SHA for a68b0f3

Commits on Mar 24, 2022

• When passing ShapedArray as input to pjit.lower() (which is not allowed normally i.e pjit(f)(*args)), consider it as Global. This will also help during the auto sharding change.

  Show description for 58efa00
  

  yashk2810

  authored and

  jax authors

  committed
  58efa00

  Copy full SHA for 58efa00

Commits on Mar 21, 2022

• Remove the duplicate _with_context_manager tests now that Mesh is the default

  Show description for 8b1cb8a
  

  yashk2810

  authored and

  jax authors

  committed
  8b1cb8a

  Copy full SHA for 8b1cb8a

Commits on Mar 19, 2022

• Set aval when creating the local_shards

  Show description for 1b17b54
  

  yashk2810

  authored and

  jax authors

  committed
  1b17b54

  Copy full SHA for 1b17b54

Commits on Mar 18, 2022

• Internal change

  Show description for 4b81311
  

  yashk2810

  authored and

  jax authors

  committed
  4b81311

  Copy full SHA for 4b81311

Commits on Mar 15, 2022

• Set in_positional_semantics should be GLOBAL for fully replicated values to avoid recompilation.

  Show description for 846c480
  

  yashk2810

  authored and

  jax authors

  committed
  846c480

  Copy full SHA for 846c480

Commits on Mar 11, 2022

• Default to normal read if the shape matches the shape the tensor was saved with.

  Show description for ee67496
  

  yashk2810

  authored and

  jax authors

  committed
  ee67496

  Copy full SHA for ee67496

Commits on Mar 9, 2022

• Add JAX support for pjit on CPU

  Show description for 6e63a90
  

  yashk2810

  authored and

  jax authors

  committed
  6e63a90

  Copy full SHA for 6e63a90

Commits on Mar 7, 2022

• Make mesh_axes on GDA strict by only allowing PartitionSpecs to be consistent with pjit.

  Show description for 99a1037
  

  yashk2810

  authored and

  jax authors

  committed
  99a1037

  Copy full SHA for 99a1037

Commits on Mar 4, 2022

• format nicely

  yashk2810

  committed
  403eae0

  Copy full SHA for 403eae0

• Format

  yashk2810

  committed
  2961071

  Copy full SHA for 2961071

Commits on Mar 3, 2022

• Fix errors

  yashk2810

  committed
  ca83227

  Copy full SHA for ca83227

• Fix errors

  yashk2810

  committed
  9d2388e

  Copy full SHA for 9d2388e

• Add pjit link

  yashk2810

  committed
  8da1397

  Copy full SHA for 8da1397

• convert to doctest

  yashk2810

  committed
  852e39e

  Copy full SHA for 852e39e

• Fix the display of GDA docs

  yashk2810

  committed
  f0df7b7

  Copy full SHA for f0df7b7

Commits on Mar 2, 2022

• Fix a typo (titled -> tiled)

  Show description for 790b4bc
  

  yashk2810

  authored and

  jax authors

  committed
  790b4bc

  Copy full SHA for 790b4bc

• Don't compute global shards if local_mesh == global_mesh and just return local_shards.

  Show description for 4755dc3
  

  yashk2810

  authored and

  jax authors

  committed
  4755dc3

  Copy full SHA for 4755dc3

• Use the new mesh property instead of the private _global_mesh attribute.

  Show description for 72cc567
  

  yashk2810

  authored and

  jax authors

  committed
  72cc567

  Copy full SHA for 72cc567

Commits on Mar 1, 2022

• Make the global_mesh that GDA has public

  Show description for 11664f8
  

  yashk2810

  authored and

  jax authors

  committed
  11664f8

  Copy full SHA for 11664f8

Commits on Feb 28, 2022

• Add block_until_ready method to GDA

  Show description for d0cc339
  

  yashk2810

  authored and

  jax authors

  committed
  d0cc339

  Copy full SHA for d0cc339

Commits on Feb 23, 2022

• Add .device() method to _DeviceArray

  Show description for a9a827e
  

  yashk2810

  authored and

  jax authors

  committed
  a9a827e

  Copy full SHA for a9a827e

• Rename unmapped_local_out_avals to out_avals since it can contain global avals (via GDA) as well as local avals.

  Show description for 98e114d
  

  yashk2810

  authored and

  jax authors

  committed
  98e114d

  Copy full SHA for 98e114d

• Fix the gpu tests that were failing with Future warning

  Show description for e2834d8
  

  yashk2810

  authored and

  jax authors

  committed
  e2834d8

  Copy full SHA for e2834d8

• Deprecate maps.mesh and replace it with maps.Mesh.

  Show description for 687a763
  

  yashk2810

  authored and

  jax authors

  committed
  687a763

  Copy full SHA for 687a763

• Make the name pjit appear in xprof and mhlo module name. Before everything was xmap even for pjit.

  Show description for e96b91d
  

  yashk2810

  authored and

  jax authors

  committed
  e96b91d

  Copy full SHA for e96b91d

Commits on Feb 22, 2022

• Make local_to_global and global_to_local private.

  Show description for afdb7b6
  

  yashk2810

  authored and

  jax authors

  committed
  afdb7b6

  Copy full SHA for afdb7b6

Commits on Feb 19, 2022

• Make resharding of GDA work if the shape is larger than what it was serialized with.

  Show description for 3290dd3
  

  yashk2810

  authored and

  jax authors

  committed
  3290dd3

  Copy full SHA for 3290dd3

• Finish jax release

  Show description for c161c62
  

  yashk2810

  authored and

  jax authors

  committed
  c161c62

  Copy full SHA for c161c62

Pagination

Previous

Next