[claude-code] Add top-level module doc for torch/distributed/tensor/_op_schema.py

ezyang · ezyang · commit dd52a1aeeca2 · 2025-07-08T22:40:00.000-04:00
Not sure how good the description is, seeking insight from maintainers. Signed-off-by: Edward Yang <ezyang@meta.com> ghstack-source-id: b6998b1 Pull-Request: #157804
diff --git a/torch/distributed/tensor/_op_schema.py b/torch/distributed/tensor/_op_schema.py
@@ -1,4 +1,27 @@
 # mypy: allow-untyped-defs
+"""
+DTensor operator schema definitions and utilities.
+
+This module defines the core data structures and utilities for describing and managing
+distributed tensor operations in PyTorch's DTensor system. It provides the foundational
+schema types used for sharding propagation, operator strategy selection, and distributed
+execution planning.
+
+Key components:
+- OpSpec: Describes acceptable sharding placements for operations
+- OpStrategy: Represents the possible sharding strategies for an operator
+- TupleStrategy: Container for multiple strategies when ops have tuple/list of tensors input
+- OpSchema: Describes operator input/output schemas with DTensorSpecs
+- OutputSharding: Manages output sharding specifications and redistribution
+- RuntimeSchemaInfo: Runtime execution metadata for operators
+- OpInfo: Complete runtime operator execution information
+
+These schema definitions enable the DTensor system to:
+1. Propagate tensor sharding information to the operator outputs
+2. Greedily select sharding strategies for distributed operations
+3. Plan and execute tensor redistributions when needed
+4. Cache sharding decisions for performance optimization
+"""
 from collections.abc import Sequence
 from dataclasses import dataclass
 from functools import cached_property
