Searched refs:ShardedAxis (Results 1 – 2 of 2) sorted by relevance
97 struct ShardedAxis { struct99 bool operator==(const ShardedAxis& other) const { return axis == other.axis; } argument100 bool operator!=(const ShardedAxis& other) const { return axis != other.axis; }112 using MeshDimAssignment = absl::variant<ShardedAxis, Replicated>;
379 py::class_<ShardedAxis> sharded_axis(pmap_lib, "ShardedAxis"); in BuildPmapSubmodule()380 sharded_axis.def(py::init<int>()).def_readonly("axis", &ShardedAxis::axis); in BuildPmapSubmodule()383 [](const ShardedAxis& x) { in BuildPmapSubmodule()386 .def("__eq__", [](const ShardedAxis& self, const ShardedAxis& other) { in BuildPmapSubmodule()