Searched refs:expanded_dims (Results 1 – 3 of 3) sorted by relevance
173 Eigen::DSizes<IndexType, NDIMS> expanded_dims;178 expanded_dims[out_index++] = in.dimension(spatial_dim);182 expanded_dims[out_index++] =188 in.reshape(merged_dims).shuffle(shuffling_perm).reshape(expanded_dims);199 Eigen::DSizes<IndexType, 4> expanded_dims;210 expanded_dims[0] = in.dimension(shuffle[0]);211 expanded_dims[1] = in.dimension(2);212 expanded_dims[2] = in.dimension(3);213 expanded_dims[3] = in.dimension(shuffle[3]);221 expanded_dims[0] = in.dimension(2);[all …]
347 HloInstruction* operand, absl::Span<const int64> expanded_dims) { in ExpandFirstDimIntoNDims() argument349 CHECK_EQ(operand->shape().dimensions(0), Product(expanded_dims)); in ExpandFirstDimIntoNDims()352 expanded_shape_dim_bounds.reserve(expanded_dims.size() + in ExpandFirstDimIntoNDims()354 absl::c_copy(expanded_dims, std::back_inserter(expanded_shape_dim_bounds)); in ExpandFirstDimIntoNDims()
185 HloInstruction* operand, absl::Span<const int64> expanded_dims);