Searched refs:Unstacked (Results 1 – 2 of 2) sorted by relevance
76 struct Unstacked { struct78 explicit Unstacked(int sz) : size(sz) {} in Unstacked() function81 bool operator==(const Unstacked& other) const { return size == other.size; } argument82 bool operator!=(const Unstacked& other) const { return size != other.size; }85 using AvalDimSharding = absl::variant<NoSharding, Chunked, Unstacked>;
51 } else if (py::isinstance<Unstacked>(value)) { in PyShardingToCpp()52 cpp_sharding.push_back(py::cast<Unstacked>(value)); in PyShardingToCpp()73 } else if (absl::holds_alternative<Unstacked>(value)) { in CppShardingToPy()74 py::handle handle = py::cast(absl::get<Unstacked>(value)); in CppShardingToPy()365 py::class_<Unstacked> unstacked(pmap_lib, "Unstacked"); in BuildPmapSubmodule()367 .def_readonly("size", &Unstacked::size) in BuildPmapSubmodule()369 [](const Unstacked& x) { in BuildPmapSubmodule()372 .def("__eq__", [](const Unstacked& self, py::object other) { in BuildPmapSubmodule()373 if (!py::isinstance<Unstacked>(other)) { in BuildPmapSubmodule()376 return self == py::cast<const Unstacked&>(other); in BuildPmapSubmodule()