Code

Use broadcasting trick for KV update#972

Code

Merged

changlan merged 4 commits intoapple:mainapple/axlearn:mainfrom changlan:fix-vlogchanglan/axlearn:fix-vlogCopy head branch name to clipboardFeb 9, 2025

Commits

Commits on Feb 7, 2025

• Use vmap and dynamic_update_slice for KV update

  changlan

  committed
  3e6f076

  Copy full SHA for 3e6f076

Commits on Feb 8, 2025

• Broadcasting trick

  changlan

  committed
  ed3ab2d

  Copy full SHA for ed3ab2d

• Simplify the impl per @markblee's suggestion

  changlan

  committed
  3d47aa0

  Copy full SHA for 3d47aa0

Commits on Feb 9, 2025

• comments

  changlan

  committed
  c86b464

  Copy full SHA for c86b464