Commits

Branch selector

 main

User selector

yashk2810

Datepicker

All time

Commit History

Commits on Dec 19, 2024

• Account for tokens in allow_spmd_sharding_propagation_to_parameters and allow_spmd_sharding_propagation_to_output compile options

  Show description for 9041b02
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  9041b02

  Copy full SHA for 9041b02

Commits on Dec 18, 2024

• [sharding_in_types] Check out_avals with mesh context too. This is because users can pass their own shardings to functions like einsum, reshape, broadcast`, etc

  Show description for af63e44
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  af63e44

  Copy full SHA for af63e44

• Remove jax_enable_memories config flag. It defaulted to True for a very long time and it's time to remove the flag.

  Show description for 8b73480
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  8b73480

  Copy full SHA for 8b73480

• [sharding_in_types] Add tests allowing inputs and outputs of jit to have different axis_types on their mesh than the axis_types on the surrounding mesh context

  Show description for 09fdd0d
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  09fdd0d

  Copy full SHA for 09fdd0d

• Allow P.UNCONSTRAINED in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs.

  Show description for e854f16
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  e854f16

  Copy full SHA for e854f16

• Fix the breakage caused by deleted enable_memories config

  Show description for 1e22149
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  1e22149

  Copy full SHA for 1e22149

• Delete enable_memories code in C++ since that flag is always True and cannot be turned off now.

  Show description for cca9afa
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  cca9afa

  Copy full SHA for cca9afa

• Don't create a mesh inside jit in sharded-computation tutorial

  Show description for 83ca7c3
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  83ca7c3

  Copy full SHA for 83ca7c3

Commits on Dec 17, 2024

• Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

  Show description for 473e2bf
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  473e2bf

  Copy full SHA for 473e2bf

• Fix save_from_both_policies in presence of save_and_offload_only_these_names by comparing the enum

  Show description for 7dd401c
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  7dd401c

  Copy full SHA for 7dd401c

Commits on Dec 13, 2024

• Set device_assignment to None when only AbstractMesh exist in the computation

  Show description for 80cf141
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  80cf141

  Copy full SHA for 80cf141

• Allow tracing and lowering (with lowering_platforms specified) to work with an AbstractMesh. Such a computation cannot be compiled.

  Show description for d0f63da
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  d0f63da

  Copy full SHA for d0f63da

• Add dynamic_arg_layouts to C++ cache and add a test in JAX which checks for cache miss if layouts of inputs arguments are different to the same jitted function.

  Show description for 1453a22
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  1453a22

  Copy full SHA for 1453a22

Commits on Dec 12, 2024

• [sharding_in_types] Make jnp.where broadcast shardings properly when a scalar exists

  Show description for 39e4f7f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  39e4f7f

  Copy full SHA for 39e4f7f

Commits on Dec 11, 2024

• [sharding_in_types] Default axis_types to Auto for all axis_names if user does not set any AxisType. Also resolve some TODOs now that we have a way for user to set the mesh.

  Show description for 41f490a
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  41f490a

  Copy full SHA for 41f490a

• [sharding_in_types] Enforce AxisTypes to always exist if set_mesh is used.

  Show description for b5e4fd1
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  b5e4fd1

  Copy full SHA for b5e4fd1

Commits on Dec 10, 2024

• Add a no-op batching rule for optimization_barrier_p

  Show description for 944d822
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  944d822

  Copy full SHA for 944d822

Commits on Dec 4, 2024

• Simply abstract_mesh and device_context context managers and handle everything via their corresponding configs in config.py

  Show description for a735bf8
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  a735bf8

  Copy full SHA for a735bf8

• [sharding_in_types] Use set_mesh API to trigger sharding_in_types instead of the config option.

  Show description for 9e2708e
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  9e2708e

  Copy full SHA for 9e2708e

• Fix the broken behavior of not resetting the abstract_mesh and device_context properly during __exit__.

  Show description for 653f654
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  653f654

  Copy full SHA for 653f654

Commits on Nov 28, 2024

• [Take 2] Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh

  Show description for 456dfeb
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  456dfeb

  Copy full SHA for 456dfeb

Commits on Nov 27, 2024

• Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh

  Show description for cc5036c
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  cc5036c

  Copy full SHA for cc5036c

• Add a private set_mesh API to enter into sharding_in_types mode. This is how users will enable sharding in types mode (with correct axis types set too but that doesn't work yet).

  Show description for 0d2dfea
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  0d2dfea

  Copy full SHA for 0d2dfea

Commits on Nov 26, 2024

• Fix a weird interaction with set_local and empty tuples passed to it.

  Show description for 6763fcf
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6763fcf

  Copy full SHA for 6763fcf

• Create a null_mesh_context internal context manager to handle null contexts properly.

  Show description for 627debc
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  627debc

  Copy full SHA for 627debc

• Add sharding argument to reshape since it also takes a shape argument for the output shape

  Show description for 59e13f8
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  59e13f8

  Copy full SHA for 59e13f8

Commits on Nov 25, 2024

• Remove _pjit_lower_cached cache. We can simplify the caching of jit as we have downstream caches and a cpp cache too.

  Show description for deab6fb
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  deab6fb

  Copy full SHA for deab6fb

• Add abstract mesh context manager to trace_context in the fallback path too (which will be deleted after jax 0.4.36 release)

  Show description for c35f8b2
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  c35f8b2

  Copy full SHA for c35f8b2

Commits on Nov 23, 2024

• When host local inputs on all hosts are the same, use _DeferredShardArg to do the transfers instead of jit to avoid blocking.

  Show description for 8699f5d
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  8699f5d

  Copy full SHA for 8699f5d

Commits on Nov 22, 2024

• [sharding_in_types] Make argmax and argmin work with sharding_in_types. This also requires adding reduce_p sharding rule

  Show description for 21f8885
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  21f8885

  Copy full SHA for 21f8885

• Use with_spec where possible to clean up the code a bit

  Show description for 7635605
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  7635605

  Copy full SHA for 7635605

• [sharding_in_types] Add scan support to sharding_in_types. There are a couple of changes here

  Show description for 355589f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  355589f

  Copy full SHA for 355589f

Commits on Nov 21, 2024

• [sharding_in_types] Add concatenate_p support

  Show description for 6568713
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6568713

  Copy full SHA for 6568713

• [sharding_in_types] Add pad_p support to sharding_in_types to handle transpose to slice correctly.

  Show description for 840cf3f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  840cf3f

  Copy full SHA for 840cf3f

Commits on Nov 20, 2024

• [sharding_in_types] Add slice_p and squeeze_p sharding rule to make flash attention work in backward pass

  Show description for 9b94180
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  9b94180

  Copy full SHA for 9b94180

Pagination

Previous

Next