Home
last modified time | relevance | path

Searched refs:num_shards_for_dim (Results 1 – 7 of 7) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Dtensor_layout.cc893 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::Layout
908 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()
H A Dtensor_layout.h317 size_t num_shards_for_dim(const ShardingSpec& dim) const;
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dbias_add_spmd_expander.cc73 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()
H A Dmeta_spmd_expander.cc199 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()
H A Dnullary_spmd_expander.cc67 op_layouts[0]->num_shards_for_dim(op_layouts[0]->dim(i)); in ExpandOp()
H A Dsoftmax_spmd_expander.cc246 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()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/
H A Dcollective_lowering.cc448 if (src_layout.num_shards_for_dim(src_layout.dim(i)) > 1 && in LowerAllGatherOp()