| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | xla_broadcast_helper_op.cc | 45 std::vector<int64> broadcast_dims; in Compile() local
47 &broadcast_dims)); in Compile()
48 if (broadcast_dims.empty()) { in Compile()
64 context, broadcast_dims.size() == min_rank_shape->dims(), in Compile()
68 absl::StrJoin(broadcast_dims, ","), "]; argument shapes: ", in Compile()
70 std::vector<int64> sorted_broadcast_dims = broadcast_dims; in Compile()
72 std::set<int64> dims_set(broadcast_dims.begin(), broadcast_dims.end()); in Compile()
74 dims_set.size() == broadcast_dims.size() && in Compile()
75 broadcast_dims == sorted_broadcast_dims, in Compile()
79 absl::StrJoin(broadcast_dims, ","), "]")); in Compile()
[all …]
|
| D | resampler_ops.cc | 60 std::vector<int64> broadcast_dims(warp_dims.begin(), warp_dims.end() - 1); in BilinearWeights() local
61 broadcast_dims.push_back(4); in BilinearWeights()
62 broadcast_dims.push_back(2); in BilinearWeights()
64 const int64 broadcast_dims_size = broadcast_dims.size(); in BilinearWeights()
70 xla::BroadcastInDim(first_term, broadcast_dims, last_two_dims_indices); in BilinearWeights()
76 std::vector<int64> ratio_broadcast_indices(broadcast_dims.size()); in BilinearWeights()
81 xla::BroadcastInDim(ratio, broadcast_dims, ratio_broadcast_indices); in BilinearWeights()
92 xla::BroadcastInDim(sign_change, broadcast_dims, last_two_dims_indices); in BilinearWeights()
|
| D | image_ops.cc | 207 std::vector<int64> broadcast_dims(input_shape.dims() - 2); in Compile() local
208 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in Compile()
209 broadcast_dims.back() = channel_dim; in Compile()
213 broadcast_dims); in Compile()
|
| /external/tensorflow/tensorflow/core/kernels/image/ |
| D | adjust_contrast_op.h | 47 Eigen::array<int, 4> broadcast_dims; in operator() local
48 broadcast_dims[0] = 1; in operator()
49 broadcast_dims[1] = height; in operator()
50 broadcast_dims[2] = width; in operator()
51 broadcast_dims[3] = 1; in operator()
61 broadcast_dims; in operator() local
62 broadcast_dims.set(1, height); in operator()
63 broadcast_dims.set(2, width); in operator()
75 .broadcast(broadcast_dims); in operator()
109 Eigen::array<int, 4> broadcast_dims; in operator() local
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | svd.cc | 137 std::vector<int64> broadcast_dims(num_dims - 1); in HouseRow() local
138 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in HouseRow()
140 auto mu = Mul(sigma, Sqrt(Square(Div(x_0j, sigma, broadcast_dims)) + one), in HouseRow()
141 broadcast_dims); in HouseRow()
145 -Mul(sigma, Div(sigma, Add(x_0j, mu), broadcast_dims), broadcast_dims)); in HouseRow()
148 (Square(Div(sigma, v_0j, broadcast_dims)) + one)); in HouseRow()
151 BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, in HouseRow()
155 beta = Select(Lt(Add(sigma, ZerosLike(beta), broadcast_dims), eps), in HouseRow()
202 std::vector<int64> broadcast_dims(num_dims - 1); in HouseCol() local
203 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in HouseCol()
[all …]
|
| D | self_adjoint_eig.cc | 165 std::vector<int64> broadcast_dims(batch_dims.size()); in Update() local
166 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in Update()
167 broadcast_dims.push_back(num_dims - 1); in Update()
178 broadcast_dims); in Update()
185 broadcast_dims); in Update()
346 std::vector<int64> broadcast_dims(num_dims - 1); in SortByEigenvalues() local
347 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in SortByEigenvalues()
348 broadcast_dims[num_dims - 2] = num_dims - 1; in SortByEigenvalues()
349 result.w = BroadcastInDim(result.w, dimensions, broadcast_dims); in SortByEigenvalues()
|
| D | arithmetic.cc | 194 std::vector<int64> broadcast_dims(input_shape.rank() - 1); in ArgMinMaxTwoPass() local
195 std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0); in ArgMinMaxTwoPass()
196 std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1); in ArgMinMaxTwoPass()
199 XlaOp select_mask = Select(Eq(input, reduced_input, broadcast_dims), in ArgMinMaxTwoPass()
208 XlaOp select_mask = Select(Eq(input, reduced_input, broadcast_dims), in ArgMinMaxTwoPass()
|
| D | svd_test.cc | 92 std::vector<int64> broadcast_dims(num_dims - 1); in ComputeMatmulUDVT() local
93 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in ComputeMatmulUDVT()
94 broadcast_dims[num_dims - 2] = num_dims - 1; in ComputeMatmulUDVT()
95 return BatchDot(Mul(u, d, broadcast_dims), TransposeInMinorDims(v), in ComputeMatmulUDVT()
|
| D | self_adjoint_eig_test.cc | 106 std::vector<int64> broadcast_dims(shape.rank() - 1); in ComputeMatmulVWVt() local
107 std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); in ComputeMatmulVWVt()
109 broadcast_dims[shape.rank() - 2] = shape.rank() - 1; in ComputeMatmulVWVt()
110 auto vw = Mul(result.v, BroadcastInDim(result.w, out_dims, broadcast_dims)); in ComputeMatmulVWVt()
|
| D | matrix.cc | 182 std::vector<int64> broadcast_dims(n_dims - 1); in SetMatrixDiagonal() local
183 absl::c_iota(broadcast_dims, 0); in SetMatrixDiagonal()
186 ++(broadcast_dims.back()); in SetMatrixDiagonal()
206 BroadcastInDim(diag, shape.dimensions(), broadcast_dims), in SetMatrixDiagonal()
243 std::vector<int64> broadcast_dims; in EinsumDiagonalLabels() local
249 broadcast_dims.push_back(dim); in EinsumDiagonalLabels()
257 return {{unique_labels, reduce_dims, broadcast_dims}}; in EinsumDiagonalLabels()
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | cwise_op_gpu_select.cu.cc | 65 Eigen::array<int, 1> broadcast_dims{size}; in operator ()() local
68 .broadcast(broadcast_dims) in operator ()()
84 Eigen::array<int, 2> broadcast_dims{{ 1, all_but_batch }}; in operator ()() local
87 Eigen::IndexList<Eigen::type2index<1>, int> broadcast_dims; in operator ()() local
88 broadcast_dims.set(1, all_but_batch); in operator ()()
103 .broadcast(broadcast_dims) in operator ()()
|
| D | cwise_op_select.cc | 341 Eigen::array<Eigen::DenseIndex, 2> broadcast_dims{{1, all_but_batch}}; in operator ()() local
344 Eigen::IndexList<Eigen::type2index<1>, Eigen::DenseIndex> broadcast_dims; in operator ()() local
345 broadcast_dims.set(1, all_but_batch); in operator ()()
352 .broadcast(broadcast_dims) in operator ()()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/lib/ |
| D | broadcast.cc | 49 std::vector<int64> broadcast_dims; in BroadcastTo() local
63 broadcast_dims.push_back(broadcast_shape.size()); in BroadcastTo()
80 absl::c_reverse(broadcast_dims); in BroadcastTo()
82 for (int64& broadcast_dim : broadcast_dims) { in BroadcastTo()
87 xla::BroadcastInDim(input, broadcast_shape, broadcast_dims); in BroadcastTo()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_helpers.cc | 88 std::vector<int64> broadcast_dims(indices_shape.dims()); in OneHot() local
89 std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0); in OneHot()
90 std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1); in OneHot()
100 xla::Eq(indices, xla::Iota(builder, iota_shape, axis), broadcast_dims), in OneHot()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/utils/ |
| D | broadcast_utils.cc | 32 DenseIntElementsAttr broadcast_dims) { in IsLegalNumpyRankedBroadcast() argument
42 if (smaller_rank != broadcast_dims.getNumElements()) { in IsLegalNumpyRankedBroadcast()
48 broadcast_dims.getIntValues().begin()); in IsLegalNumpyRankedBroadcast()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/python/ |
| D | xla.py | 122 def broadcasting_binary_op_wrapper(x, y, broadcast_dims=None, name=None): argument
124 broadcast_dims = broadcast_dims or []
125 broadcast_dims = ops.convert_to_tensor(broadcast_dims, dtypes.int64)
129 x, y = gen_xla_ops.xla_broadcast_helper(x, y, broadcast_dims)
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/utils/ |
| D | broadcast_utils.h | 37 DenseIntElementsAttr broadcast_dims);
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
| D | legalize_to_linalg.cc | 547 auto broadcast_dims = std::get<1>(operand_and_dims); in matchAndRewrite() local
558 result_type.getDimSize(broadcast_dims.front())) { in matchAndRewrite()
572 auto indexing_maps = getIndexingMaps(op, broadcast_dims, result_shape, in matchAndRewrite()
606 SmallVector<int64_t, 2> new_shape, new_strides, broadcast_dims; in InsertReshapeIfNecessary() local
623 broadcast_dims.push_back(dim); in InsertReshapeIfNecessary()
636 broadcast_dims.push_back(0); in InsertReshapeIfNecessary()
653 return std::make_pair(operand, broadcast_dims); in InsertReshapeIfNecessary()
657 ArrayRef<int64_t> broadcast_dims, in getIndexingMaps() argument
673 for (const auto& broadcast_dim : llvm::enumerate(broadcast_dims)) { in getIndexingMaps()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | convolution_group_converter.cc | 462 std::vector<int64> broadcast_dims; in HandleConvolution() local
467 broadcast_dims.push_back(i); in HandleConvolution()
470 expanded_filter_shape, reshaped_filter, broadcast_dims)); in HandleConvolution()
|
| D | indexed_array_analysis.cc | 876 absl::Span<const int64> broadcast_dims = broadcast_instr->dimensions(); in ComputeArrayForElementwiseBinaryOp() local
878 return absl::c_find(broadcast_dims, output_dim) == broadcast_dims.end(); in ComputeArrayForElementwiseBinaryOp()
893 for (int64 broadcast_dim : broadcast_dims) { in ComputeArrayForElementwiseBinaryOp()
|
| D | space_to_batch_converter.cc | 1034 std::vector<int64> broadcast_dims; in PropagateOnBroadcast() local
1036 broadcast_dims.push_back(DimLookUp(permute_dims, j)); in PropagateOnBroadcast()
1039 broadcast_dims, final_shape_dims); in PropagateOnBroadcast()
1107 auto broadcast_dims = broadcast->dimensions(); in IsBroadcastPropagatable() local
1108 return !absl::c_linear_search(broadcast_dims, space_dim); in IsBroadcastPropagatable()
|
| D | algebraic_simplifier.cc | 1355 absl::InlinedVector<int64, 8> broadcast_dims; in HandleConcatenate() local
1360 broadcast_dims.push_back(i); in HandleConcatenate()
1366 broadcast_dims, concatenate->shape())); in HandleConcatenate()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 425 DenseIntElementsAttr broadcast_dims, in StaticBinaryBroadcast() argument
429 auto result_type = GetStaticBroadcastType(x_type, y_type, broadcast_dims); in StaticBinaryBroadcast()
432 << " with broadcast_dims = " << broadcast_dims; in StaticBinaryBroadcast()
439 x = builder.create<BroadcastInDimOp>(loc, result_type, x, broadcast_dims); in StaticBinaryBroadcast()
451 y = builder.create<BroadcastInDimOp>(loc, result_type, y, broadcast_dims); in StaticBinaryBroadcast()
508 auto broadcast_dims = GetI64ElementsAttr({feature_dim}, &builder); in Broadcast1DToFeatureDim() local
515 loc, to_type, broadcast_from, result_extents, broadcast_dims); in Broadcast1DToFeatureDim()
533 auto broadcast_dims = GetI64ElementsAttrForSeq(0, rank, &builder); in BroadcastToShapeOf() local
535 loc, to_type, input, result_extents, broadcast_dims); in BroadcastToShapeOf()
3685 auto broadcast_dims = GetI64ElementsAttr({}, &rewriter); in matchAndRewrite() local
[all …]
|
| /external/tensorflow/tensorflow/compiler/tests/ |
| D | xla_ops_test.py | 66 lambda x, y: xla.add(x, y, broadcast_dims=(0,)),
72 lambda x, y: xla.add(x, y, broadcast_dims=(1,)),
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 231 absl::flat_hash_set<int64> broadcast_dims; in AddInstruction() local
233 broadcast_dims.insert(i); in AddInstruction()
238 broadcast_dims.clear(); in AddInstruction()
243 broadcast_dims.erase(i); in AddInstruction()
247 if (!broadcast_dims.empty()) { in AddInstruction()
248 broadcast_dims_[hlo] = std::move(broadcast_dims); in AddInstruction()
1317 if (auto broadcast_dims = state_.b->BroadcastDimsForCreatedHlo(hlo())) { in ReshardWithCollectivePermute() local
1318 if (!(*broadcast_dims)->empty()) { in ReshardWithCollectivePermute()
1323 if ((*broadcast_dims)->contains(i)) { in ReshardWithCollectivePermute()
|