Home
last modified time | relevance | path

Searched refs:Unstacked (Results 1 – 2 of 2) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/python/
Dpmap_lib.h76 struct Unstacked { struct
78 explicit Unstacked(int sz) : size(sz) {} in Unstacked() argument
81 bool operator==(const Unstacked& other) const { return size == other.size; } argument
82 bool operator!=(const Unstacked& other) const { return size != other.size; }
85 using AvalDimSharding = absl::variant<NoSharding, Chunked, Unstacked>;
Dpmap_lib.cc51 } 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()