Searched refs:ShardingSpec (Results 1 – 2 of 2) sorted by relevance
116 class ShardingSpec {118 ShardingSpec(std::vector<AvalDimSharding> sharding, in ShardingSpec() function158 pybind11::handle aval, ShardingSpec sharding_spec, in ShardedDeviceArray()170 const ShardingSpec& GetShardingSpec() const { return sharding_spec_; } in GetShardingSpec()177 ShardingSpec sharding_spec_;
189 std::vector<ShardingSpec> sharding_specs_;401 py::class_<ShardingSpec> sharding_spec(pmap_lib, "ShardingSpec"); in BuildPmapSubmodule()406 .def_property_readonly("sharding", &ShardingSpec::GetSharding) in BuildPmapSubmodule()407 .def_property_readonly("mesh_mapping", &ShardingSpec::GetMeshMapping); in BuildPmapSubmodule()410 sda.def(py::init<pybind11::handle, ShardingSpec, pybind11::list>()) in BuildPmapSubmodule()