Home
last modified time | relevance | path

Searched refs:broadcast_dim (Results 1 – 10 of 10) sorted by relevance

/external/tensorflow/tensorflow/lite/kernels/internal/reference/
Dbatch_matmul.h31 inline int broadcast_dim(int lhs_dim, int rhs_dim) { in broadcast_dim() function
61 const int batch_dim0 = batch_matmul::broadcast_dim( in BatchMatMul()
63 const int batch_dim1 = batch_matmul::broadcast_dim( in BatchMatMul()
65 const int batch_dim2 = batch_matmul::broadcast_dim( in BatchMatMul()
119 const int batch_dim0 = batch_matmul::broadcast_dim( in BatchMatMul()
121 const int batch_dim1 = batch_matmul::broadcast_dim( in BatchMatMul()
123 const int batch_dim2 = batch_matmul::broadcast_dim( in BatchMatMul()
208 const int batch_dim0 = batch_matmul::broadcast_dim( in BatchMatMul()
210 const int batch_dim1 = batch_matmul::broadcast_dim( in BatchMatMul()
212 const int batch_dim2 = batch_matmul::broadcast_dim( in BatchMatMul()
/external/tensorflow/tensorflow/lite/kernels/internal/optimized/
Dbatch_matmul.h41 auto broadcast_dim = [](int lhs_dim, int rhs_dim) { in BatchMatMul()
62 broadcast_dim(extended_lhs_shape.Dims(0), extended_rhs_shape.Dims(0)); in BatchMatMul()
64 broadcast_dim(extended_lhs_shape.Dims(1), extended_rhs_shape.Dims(1)); in BatchMatMul()
66 broadcast_dim(extended_lhs_shape.Dims(2), extended_rhs_shape.Dims(2)); in BatchMatMul()
132 auto broadcast_dim = [](int lhs_dim, int rhs_dim) { in BatchMatMul()
153 broadcast_dim(extended_lhs_shape.Dims(0), extended_rhs_shape.Dims(0)); in BatchMatMul()
155 broadcast_dim(extended_lhs_shape.Dims(1), extended_rhs_shape.Dims(1)); in BatchMatMul()
157 broadcast_dim(extended_lhs_shape.Dims(2), extended_rhs_shape.Dims(2)); in BatchMatMul()
286 auto broadcast_dim = [](int lhs_dim, int rhs_dim) { in BatchMatMul()
307 broadcast_dim(extended_lhs_shape.Dims(0), extended_rhs_shape.Dims(0)); in BatchMatMul()
[all …]
/external/tensorflow/tensorflow/compiler/tf2xla/lib/
Dbroadcast.cc82 for (int64& broadcast_dim : broadcast_dims) { in BroadcastTo()
83 broadcast_dim = broadcast_shape_size - broadcast_dim - 1; in BroadcastTo()
/external/tensorflow/tensorflow/lite/kernels/
Dbatch_matmul.cc109 int broadcast_dim = lhs_dim; in ResizeOutputTensor() local
111 broadcast_dim = rhs_dim; in ResizeOutputTensor()
113 output_shape->data[i] = broadcast_dim; in ResizeOutputTensor()
/external/tensorflow/tensorflow/compiler/xla/service/
Dindexed_array_analysis.cc893 for (int64 broadcast_dim : broadcast_dims) { in ComputeArrayForElementwiseBinaryOp() local
894 simulated_index[broadcast_dim] = IndexComponent::NotBroadcasted; in ComputeArrayForElementwiseBinaryOp()
Ddynamic_dimension_inference.cc239 int64 broadcast_dim = hlo->dimensions(dimension); in HandleBroadcast() local
240 parent_->SetDynamicSize(hlo, {}, broadcast_dim, dynamic_size); in HandleBroadcast()
Ddynamic_padder.cc1374 HloInstruction* broadcast_dim = comp->AddInstruction( in RewriteDynamicUpdateSlice() local
1377 mask_shape_pred, iota, broadcast_dim, ComparisonDirection::kLt)); in RewriteDynamicUpdateSlice()
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/
Dlegalize_to_linalg.cc673 for (const auto& broadcast_dim : llvm::enumerate(broadcast_dims)) { in getIndexingMaps() local
674 int size = broadcast_dim.value(); in getIndexingMaps()
676 operand_shape[broadcast_dim.index()] == 1 && result_shape[size] != 1; in getIndexingMaps()
/external/tensorflow/tensorflow/compiler/xla/tests/
Darray_elementwise_ops_test.cc2032 for (int broadcast_dim : {0, 1}) { in XLA_TEST_F()
2036 Max(u, v, /*broadcast_dimensions=*/{broadcast_dim}); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf.cc4008 int64_t broadcast_dim; in matchAndRewrite() local
4012 broadcast_dim = broadcasted_shape.size(); in matchAndRewrite()
4017 broadcast_dim = broadcasted_shape.size(); in matchAndRewrite()
4020 broadcast_dimensions.push_back(broadcast_dim); in matchAndRewrite()