Pallas 设计笔记

设计

• Pallas 设计
  • 引言
  • Pallas：为内核扩展 JAX
• Pallas 异步操作
  • 背景 + 动机
  • 如何在 TPU 上实现分解的 Pallas 异步操作？
  • 为什么不起作用？
  • 状态的复仇
  • 结论