Searched refs:num_shards_for_dim (Results 1 – 7 of 7) sorted by relevance
893 num_shards.push_back(num_shards_for_dim(sharding_spec)); in num_shards()898 size_t Layout::num_shards_for_dim(const ShardingSpec& dim) const { in num_shards_for_dim() function in tensorflow::dtensor::Layout908 if (num_shards_for_dim(sharding_spec) > 1) { in IsFullyReplicated()917 (num_shards_for_dim(sharding_specs_.back()) == 1); in IsLastDimReplicated()927 if (num_shards_for_dim(dim) != 1) { in IsBatchParallel()938 if (num_shards_for_dim(sharding_specs_[i]) != 1) return false; in IsBatchParallel()958 if ((this->num_shards_for_dim(this->sharding_specs_[i]) != 1) || in IsEquivalent()959 (b.num_shards_for_dim(b.sharding_specs_[i]) != 1)) in IsEquivalent()
317 size_t num_shards_for_dim(const ShardingSpec& dim) const;
73 input_layout.num_shards_for_dim(input_layout.dim(c_dim_idx)); in ExpandOp()75 output_layout.num_shards_for_dim(output_layout.dim(c_dim_idx)); in ExpandOp()
199 if (input_layout->num_shards_for_dim(input_layout->dim(axis)) != 1) { in ExpandOp()260 layout.num_shards_for_dim(tensor_dimension) > 1) { in VerifyPaddedDimensionNotSharded()352 operand_layout.num_shards_for_dim(tensor_dimension) > 1) in VerifyTileOperandLayout()648 const int num_shards = input_layout.num_shards_for_dim( in MakeLayoutForReshape()731 input_layout->num_shards_for_dim(input_layout->dim(input_start)); in ExpandOp()
67 op_layouts[0]->num_shards_for_dim(op_layouts[0]->dim(i)); in ExpandOp()
246 desired_layout.num_shards_for_dim(desired_layout.sharding_specs().back()); in ComputeOneHot()248 int64_t num_shards = desired_layout.num_shards_for_dim(desired_layout.dim(1)); in ComputeOneHot()
448 if (src_layout.num_shards_for_dim(src_layout.dim(i)) > 1 && in LowerAllGatherOp()