tenzr::NamedSharding¶

Pairs a device Mesh with a PartitionSpec to describe a tensor layout.

#include <named_sharding.h>

Public Functions¶

Public Functions Documentation¶

function NamedSharding¶

NamedSharding() =default

function NamedSharding¶

inline NamedSharding(
    Mesh mesh,
    PartitionSpec spec
)

function mesh¶

inline const Mesh & mesh() const

The device mesh.

function Spec¶

inline const PartitionSpec & Spec() const

The partition spec.

function ToString¶

inline std::string ToString() const

Render the sharding as a string.

function ToShardyString¶

inline std::string ToShardyString(
    std::string_view mesh_sym ="mesh"
) const

Generate Shardy MLIR attribute string: #sdy.sharding<@mesh, [{"data"}, {}]>

function ToShardyConstraint¶

inline std::string ToShardyConstraint(
    std::string_view out_var,
    std::string_view in_var,
    std::string_view tensor_type,
    std::string_view mesh_sym ="mesh"
) const

Generate Shardy MLIR sharding constraint operation: out = sdy.sharding_constraint in <@mesh, [...]> : tensor<...>

function ToGspmdString¶

inline std::string ToGspmdString(
    const Shape & tensor_shape
) const

Lower into OpenXLA / GSPMD device tile notation string: e.g. {devices=[4,2]<=8 last_tile_dim_replicate}.

function operator==¶

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

function operator!=¶

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