tenzr::PartitionSpec¶

Per-dimension list of mesh axis names describing how a tensor is partitioned.

#include <partition_spec.h>

Public Types¶

Public Functions¶

Public Attributes¶

Public Types Documentation¶

using DimSpec¶

using tenzr::PartitionSpec::DimSpec =  std::vector<std::string>;

Public Functions Documentation¶

function PartitionSpec¶

PartitionSpec() =default

function PartitionSpec¶

inline PartitionSpec(
    std::initializer_list< DimSpec > dims
)

function PartitionSpec¶

inline explicit PartitionSpec(
    std::vector< DimSpec > dims
)

function PartitionSpec¶

inline explicit PartitionSpec(
    std::initializer_list< std::string > single_axes
)

function Rank¶

inline size_t Rank() const

Number of tensor dimensions described.

function Empty¶

inline bool Empty() const

Whether the spec describes no dimensions.

function Partitions¶

inline const std::vector< DimSpec > & Partitions() const

Per-dimension axis-name lists.

function operator[]¶

inline const DimSpec & operator[](
    size_t index
) const

function IsReplicated¶

inline bool IsReplicated(
    size_t index
) const

Parameters:

• index Tensor dimension index.

Return: True when the dimension is unpartitioned.

Whether a dimension is replicated.

function IsUnconstrained¶

inline bool IsUnconstrained(
    size_t index
) const

Parameters:

• index Tensor dimension index.

Return: True when the dimension carries the ? sentinel.

Whether a dimension is unconstrained.

function ToString¶

inline std::string ToString() const

Render the spec as a string.

function ToShardyDimsString¶

inline std::string ToShardyDimsString() const

Format partitions for Shardy MLIR attribute: e.g. [{"data"}, {}] or [{"data"}, {?}].

function operator==¶

inline bool operator==(
    const PartitionSpec & other
) const

function operator!=¶

inline bool operator!=(
    const PartitionSpec & other
) const

Public Attributes Documentation¶

variable kNone¶

static const DimSpec kNone {};

Sentinel representing an unpartitioned / replicated dimension.

variable kUnconstrained¶

static const DimSpec kUnconstrained {"?"};

Sentinel representing an unconstrained dimension in Shardy.