Searched refs:dst_dim (Results 1 – 3 of 3) sorted by relevance
165 for (auto dst_dim : dst_sparse_dims) { in DenseToSparse() local166 if (!dst_dim_has_nonzeroes[dst_dim]) { in DenseToSparse()169 dim_metadata_[2 * dst_dim + 1].push_back(coordinate[dst_dim]); in DenseToSparse()170 dst_dim_has_nonzeroes[dst_dim] = true; in DenseToSparse()
1214 const int dst_dim = IsColMajor ? i : NumDims - i - 1;1215 if (dst.dims[dst_dim] != 1) break;1240 const int dst_dim = IsColMajor ? i : NumDims - i - 1;1241 const IndexType dst_stride = dst.strides[dst_dim];1242 const IndexType src_stride = src.strides[dim_map[dst_dim]];1244 dst_inner_dim_size *= dst.dims[dst_dim];1264 const int dst_dim = IsColMajor ? i + 1 : NumDims - i - 2;1265 if (dst.dims[dst_dim] == 1) continue;1267 it[idx].size = dst.dims[dst_dim];1268 it[idx].input_stride = src.strides[dim_map[dst_dim]];[all …]
123 int64_t dst_dim = dst->tile_assignment().dim(i); in MergeShardingIfCompatible() local125 if (dst_dim == 1) { in MergeShardingIfCompatible()128 merged_tile_dims.push_back(dst_dim); in MergeShardingIfCompatible()129 } else if (dst_dim == merge_dim) { in MergeShardingIfCompatible()130 merged_tile_dims.push_back(dst_dim); in MergeShardingIfCompatible()