| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | dot_decomposer.cc | 42 const int64 num_batch_dims = original_dnums.lhs_batch_dimensions_size(); in CanonicalizeDot() local
49 lhs_rank - num_batch_dims - num_contracting_dims; in CanonicalizeDot()
56 batch_dim_sizes.reserve(num_batch_dims); in CanonicalizeDot()
99 rhs_rank - num_batch_dims - num_contracting_dims; in CanonicalizeDot()
152 for (int64 i = 0; i < num_batch_dims; ++i) { in CanonicalizeDot()
157 num_batch_dims + (lhs_non_contracting_size > 1 ? 1 : 0)); in CanonicalizeDot()
158 dot_dnums.add_rhs_contracting_dimensions(num_batch_dims); in CanonicalizeDot()
|
| D | qr_expander.cc | 189 const int64 num_batch_dims = num_dims - 2; in QrBlock() local
190 std::vector<int64> batch_dims(num_batch_dims); in QrBlock()
191 for (int i = 0; i < num_batch_dims; ++i) { in QrBlock()
195 std::vector<int64> batch_dim_indices(num_batch_dims); in QrBlock()
369 const int64 num_batch_dims = num_dims - 2; in BuildQrDecomposition() local
370 std::vector<int64> batch_dims(num_batch_dims); in BuildQrDecomposition()
371 for (int i = 0; i < num_batch_dims; ++i) { in BuildQrDecomposition()
|
| D | transpose_folding.cc | 47 int64 num_batch_dims = in CanFoldOperandsIntoDot() local
49 int64 expected_rank = 2 + num_batch_dims; in CanFoldOperandsIntoDot()
|
| D | elemental_ir_emitter.cc | 2220 int64 num_batch_dims = dim_numbers.rhs_batch_dimensions_size(); in EmitElementalDot() local
2221 for (int64 i = 0; i < num_batch_dims; i++) { in EmitElementalDot()
2225 for (int64 i = 0; i < rhs_dims - 1 - num_batch_dims; i++) { in EmitElementalDot()
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | svd.cc | 118 const int64 num_batch_dims = num_dims - 2; in HouseRow() local
119 std::vector<int64> batch_dims(num_batch_dims); in HouseRow()
120 for (int k = 0; k < num_batch_dims; ++k) { in HouseRow()
183 const int64 num_batch_dims = num_dims - 2; in HouseCol() local
184 std::vector<int64> batch_dims(num_batch_dims); in HouseCol()
185 for (int k = 0; k < num_batch_dims; ++k) { in HouseCol()
257 const int64 num_batch_dims = num_dims - 2; in HouseHolderBidiagonalization() local
258 std::vector<int64> batch_dims(num_batch_dims); in HouseHolderBidiagonalization()
259 for (int i = 0; i < num_batch_dims; ++i) { in HouseHolderBidiagonalization()
459 const int64 num_batch_dims = num_dims - 2; in OneSidedJacobiUpdate() local
[all …]
|
| D | lu_decomposition.cc | 34 const int num_batch_dims = a_shape.dimensions().size() - 2; in LuDecomposition() local
37 a_shape.dimensions().begin() + num_batch_dims); in LuDecomposition()
|
| D | self_adjoint_eig.cc | 429 const int64 num_batch_dims = num_dims - 2; in SelfAdjointEig() local
430 std::vector<int64> batch_dims(num_batch_dims); in SelfAdjointEig()
431 for (int i = 0; i < num_batch_dims; ++i) { in SelfAdjointEig()
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | nn_ops.py | 194 num_batch_dims=1): argument
199 input_shape.ndims - num_batch_dims + 1)
206 filter_shape.ndims + num_batch_dims - 1)
211 if input_shape.ndims < 3 or input_shape.ndims - num_batch_dims + 1 > 5:
215 .format(input_shape.ndims, num_batch_dims))
216 conv_dims = input_shape.ndims - num_batch_dims - 1
601 num_batch_dims=1): argument
616 starting_spatial_dim = num_batch_dims + 1
618 starting_spatial_dim = num_batch_dims
1093 num_batch_dims = inputs_rank - num_spatial_dims - 1 # Channel dimension.
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | dot_op_emitter.cc | 1134 int64 num_batch_dims = in EmitBatchDotOperation() local
1143 CollapseFirstNDims(b, lhs_array, num_batch_dims); in EmitBatchDotOperation()
1145 CollapseFirstNDims(b, rhs_array, num_batch_dims); in EmitBatchDotOperation()
1147 CollapseFirstNDims(b, target_array, num_batch_dims); in EmitBatchDotOperation()
1171 dot_info.dim_nums.lhs_contracting_dimensions(0) - num_batch_dims); in EmitBatchDotOperation()
1174 dot_info.dim_nums.rhs_contracting_dimensions(0) - num_batch_dims); in EmitBatchDotOperation()
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | conv_ops_3d_test.py | 229 self.assertEqual(convolver1.num_batch_dims, 1)
235 self.assertEqual(convolver2.num_batch_dims, 2)
|
| D | conv_ops_test.py | 472 self.assertEqual(convolver1.num_batch_dims, 1)
478 self.assertEqual(convolver2.num_batch_dims, 2)
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 540 bool transpose_rhs, int64_t num_batch_dims, in BatchDot() argument
543 llvm::to_vector<4>(llvm::seq<int64_t>(0, num_batch_dims)), builder); in BatchDot()
545 llvm::makeArrayRef({transpose_lhs ? num_batch_dims : num_batch_dims + 1}), in BatchDot()
548 llvm::makeArrayRef({transpose_rhs ? num_batch_dims + 1 : num_batch_dims}), in BatchDot()
5825 const int64_t num_batch_dims = num_dims - 2; in QRBlock() local
5826 auto batch_dims = a_type.getShape().take_front(num_batch_dims); in QRBlock()
5857 auto vva = BatchDot(loc, v_broadcast, false, a, false, num_batch_dims, in QRBlock()
5859 vva = BatchDot(loc, v_broadcast, true, vva, false, num_batch_dims, in QRBlock()
5888 GetI64ElementsAttr(llvm::SmallVector<int64_t, 4>(num_batch_dims, 1), in QRBlock()
5903 const int64_t minor_dim = num_batch_dims; in QRBlock()
[all …]
|
| /external/llvm-project/mlir/include/mlir/Dialect/Vector/ |
| D | VectorOps.td | 60 num_batch_dims (see dimension type descriptions below)). For K = 0 (no
|