// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <sstream>
#include "bindings.h"
namespace tenzr::python {
void BindSharding(nb::module_& m) {
auto shard_m = m.def_submodule("sharding", "Multi-device mesh, PartitionSpec, and NamedSharding with OpenXLA Shardy support");
// Mesh
nb::class_<Mesh>(shard_m, "Mesh")
.def(
"__init__",
[](Mesh* self, nb::handle devices_or_shape, std::vector<std::string> axis_names) {
if (nb::isinstance<Array>(devices_or_shape)) {
Array arr = nb::cast<Array>(devices_or_shape);
std::vector<int64_t> devs = arr.ToVectorAs<int64_t>();
std::vector<int64_t> shape(arr.shape().Rank());
for (size_t i = 0; i < arr.shape().Rank(); ++i) {
shape[i] = arr.shape().Dim(i);
}
new (self) Mesh(devs, shape, axis_names);
} else {
auto shape = nb::cast<std::vector<int64_t>>(devices_or_shape);
new (self) Mesh(shape, axis_names);
}
},
"devices_or_shape"_a, "axis_names"_a)
.def_prop_ro("devices", &Mesh::Devices)
.def_prop_ro("axis_names", &Mesh::AxisNames)
.def_prop_ro("num_devices", &Mesh::NumDevices)
.def_prop_ro("shape", &Mesh::shape)
.def("to_shardy_mesh_decl", &Mesh::ToShardyMeshDecl, "mesh_sym"_a = "mesh")
.def("__repr__", &Mesh::ToString);
// PartitionSpec
nb::class_<PartitionSpec>(shard_m, "PartitionSpec")
.def("__init__",
[](PartitionSpec* self, nb::args args) {
std::vector<PartitionSpec::DimSpec> dims;
dims.reserve(args.size());
for (size_t i = 0; i < args.size(); ++i) {
auto item = args[i];
if (item.is_none()) {
dims.push_back({});
} else if (nb::isinstance<nb::str>(item)) {
dims.push_back({nb::cast<std::string>(item)});
} else if (nb::isinstance<nb::tuple>(item) || nb::isinstance<nb::list>(item)) {
dims.push_back(nb::cast<std::vector<std::string>>(item));
} else {
dims.push_back({});
}
}
new (self) PartitionSpec(std::move(dims));
})
.def_prop_ro("rank", &PartitionSpec::Rank)
.def_prop_ro("partitions", &PartitionSpec::Partitions)
.def("__repr__", &PartitionSpec::ToString);
// NamedSharding
nb::class_<NamedSharding>(shard_m, "NamedSharding")
.def(nb::init<Mesh, PartitionSpec>(), "mesh"_a, "spec"_a)
.def_prop_ro("mesh", &NamedSharding::mesh)
.def_prop_ro("spec", &NamedSharding::Spec)
.def(
"to_gspmd_string", [](const NamedSharding& ns, std::vector<int64_t> shape) { return ns.ToGspmdString(Shape(shape)); }, "shape"_a)
.def("to_shardy_string", &NamedSharding::ToShardyString, "mesh_sym"_a = "mesh")
.def("__repr__", &NamedSharding::ToString);
// ShardingAnalysis
nb::class_<sharding::ShardingAnalysis>(shard_m, "ShardingAnalysis")
.def_prop_ro("mesh_decl", [](const sharding::ShardingAnalysis& a) { return a.mesh_decl; })
.def_prop_ro("shardy_attr", [](const sharding::ShardingAnalysis& a) { return a.shardy_attr; })
.def_prop_ro("gspmd_attr", [](const sharding::ShardingAnalysis& a) { return a.gspmd_attr; })
.def_prop_ro("shardy_constraint_op", [](const sharding::ShardingAnalysis& a) { return a.shardy_constraint_op; })
.def_prop_ro("is_replicated", [](const sharding::ShardingAnalysis& a) { return a.is_replicated; })
.def_prop_ro("is_unconstrained", [](const sharding::ShardingAnalysis& a) { return a.is_unconstrained; })
.def("format_report", &sharding::ShardingAnalysis::FormatReport);
shard_m.def(
"analyze_sharding",
[](const NamedSharding& ns, std::vector<int64_t> shape) -> sharding::ShardingAnalysis { return sharding::AnalyzeSharding(ns, Shape(shape)); },
"sharding"_a, "shape"_a);
shard_m.def(
"format_shardy_module_header",
[](const Mesh& mesh, std::string func_name, const std::vector<NamedSharding>& arg_shardings, const std::vector<std::vector<int64_t>>& arg_shapes,
std::string mesh_sym) -> std::string {
std::vector<Shape> shapes;
shapes.reserve(arg_shapes.size());
for (const auto& s : arg_shapes) {
shapes.emplace_back(s);
}
return sharding::FormatShardyModuleHeader(mesh, func_name, arg_shardings, shapes, mesh_sym);
},
"mesh"_a, "func_name"_a = "main", "arg_shardings"_a = std::vector<NamedSharding>{}, "arg_shapes"_a = std::vector<std::vector<int64_t>>{},
"mesh_sym"_a = "mesh");
}
} // namespace tenzr::python