tenzr::ShardingSpec¶

More...

#include <sharding.h>

Public Functions¶

Detailed Description¶

class tenzr::ShardingSpec;

Partitioning of an array across devices, either replicated or tiled over a device mesh.

Public Functions Documentation¶

function ShardingSpec¶

inline ShardingSpec()

function ShardingSpec¶

explicit ShardingSpec(
    const NamedSharding & ns
)

Parameters:

• ns Named sharding to translate.

Build a ShardingSpec from a NamedSharding, marking it tiled when any tensor dimension is partitioned.

function type¶

inline ShardingType type() const

The sharding kind.

function IsReplicated¶

inline bool IsReplicated() const

Whether the spec is replicated.

function IsTiled¶

inline bool IsTiled() const

Whether the spec is tiled.

function MeshShape¶

inline const std::vector< int64_t > & MeshShape() const

The device mesh shape.

function DimToAxis¶

inline const std::vector< int > & DimToAxis() const

Mapping from tensor dimension to mesh axis.

function ToString¶

std::string ToString() const

Human-readable description.

Return: Textual description of the sharding.

Render the sharding spec as a human-readable string.

function ToShardingAttribute¶

std::string ToShardingAttribute() const

Lower to a StableHLO/Shardy sharding attribute string.

Return: Sharding attribute such as {replicated} or {devices=[...]}.

Emit the XLA sharding attribute string for this spec.

function Replicated¶

static ShardingSpec Replicated()

A replicated (no partitioning) spec.

Return: A Replicated spec.

Create a fully replicated sharding spec.

function Tiled¶

static ShardingSpec Tiled(
    std::vector< int64_t > mesh_shape,
    std::vector< int > tensor_dim_to_mesh_axis
)

A tiled spec with mesh_shape and tensor-to-mesh axis mapping tensor_dim_to_mesh_axis.

Parameters:

• mesh_shape Dimensions of the device mesh.
• tensor_dim_to_mesh_axis Mesh axis assigned to each tensor dimension; -1 for replicated.

Return: A Tiled spec.

Create a tiled sharding spec from a mesh shape and tensor-dim to mesh-axis mapping.