Exploring how optimizations for GEMMs work

This module defines a type system for distributed training code, based off of JAX's sharding in types, but adapted for the PyTorch ecosystem.

Local and global SPMD type system supporting erasure

Optimized primitives for collective multi-GPU communication