Einshard Documentation

Useful links: Source Repository | Issue Tracker

Welcome to einshard’s documentation! Einshard is an einsum-like high-level API to shard and replicate arrays in JAX. With einshard, you can easily implement data parallelism and 1-D tensor parallelism.

• Quick Start
  • Installation
  • Shard a JAX Array Using Einshard
• Syntax
  • Einshard Expression
  • Sharding and Replication Methods
• Tutorial
  • Concepts: Axes and Parallelism
  • Steps for Parallelizing a Model Using Einshard
  • Example: 1-D Tensor Parallelism for a 2-Layer MLP
  • General Approach for Larger Models
• Theory
• API References
  • einshard()
  • make_sharding()
• Development Guide