Searched refs:sharding_type (Results 1 – 4 of 4) sorted by relevance
448 handle sharding_type = getattr(handle_obj, "type");449 if (!sharding_type.is_none()) {450 value.set_type(sharding_type.cast<xla::OpSharding_Type>());476 handle sharding_type = getattr(tuple_sharding, "type");477 if (!sharding_type.is_none()) {478 sharding->set_type(sharding_type.cast<xla::OpSharding_Type>());
630 const auto sharding_type = sharding.type(); in GetMetadataArgumentMapping() local631 if (sharding_type == xla::OpSharding::OTHER) { in GetMetadataArgumentMapping()634 } else if (sharding_type == xla::OpSharding::REPLICATED) { in GetMetadataArgumentMapping()637 assert(sharding_type == xla::OpSharding::MAXIMAL); in GetMetadataArgumentMapping()
273 template <typename StorageIndex, int sharding_type>275 sharding_type> {298 if (sharding_type == ShardByCol) {341 template <typename StorageIndex, int sharding_type>343 StorageIndex, sharding_type> {
316 TpuProgramShardingType sharding_type) const { in tpu_programs()322 tpu_programs_[i], sharding_type)); in tpu_programs()