Searched refs:num_batch_dims (Results 1 – 6 of 6) sorted by relevance
46 const int64 num_batch_dims = dnums.lhs_batch_dimensions_size(); in DecomposeBatchDot() local50 for (int i = 0; i < num_batch_dims; ++i) { in DecomposeBatchDot()59 const bool lhs_transpose = (lhs_contracting_dim_number - num_batch_dims) == 0; in DecomposeBatchDot()63 const bool rhs_transpose = (rhs_contracting_dim_number - num_batch_dims) == 1; in DecomposeBatchDot()67 const int64 lhs_rows = lhs_shape.dimensions(num_batch_dims + 0); in DecomposeBatchDot()68 const int64 lhs_cols = lhs_shape.dimensions(num_batch_dims + 1); in DecomposeBatchDot()78 const int64 rhs_rows = rhs_shape.dimensions(num_batch_dims + 0); in DecomposeBatchDot()79 const int64 rhs_cols = rhs_shape.dimensions(num_batch_dims + 1); in DecomposeBatchDot()89 const int64 dot_rows = dot_shape.dimensions(num_batch_dims + 0); in DecomposeBatchDot()90 const int64 dot_cols = dot_shape.dimensions(num_batch_dims + 1); in DecomposeBatchDot()[all …]
2161 int64 num_batch_dims = dim_numbers.rhs_batch_dimensions_size(); in EmitElementalDot() local2162 for (int64 i = 0; i < num_batch_dims; i++) { in EmitElementalDot()2166 for (int64 i = 0; i < rhs_dims - 1 - num_batch_dims; i++) { in EmitElementalDot()
167 const int64 num_batch_dims = num_dims - 2; in QRBlock() local168 std::vector<int64> batch_dims(num_batch_dims); in QRBlock()169 for (int i = 0; i < num_batch_dims; ++i) { in QRBlock()173 std::vector<int64> batch_dim_indices(num_batch_dims); in QRBlock()344 const int64 num_batch_dims = num_dims - 2; in QRDecomposition() local345 std::vector<int64> batch_dims(num_batch_dims); in QRDecomposition()346 for (int i = 0; i < num_batch_dims; ++i) { in QRDecomposition()
123 const int64 num_batch_dims = num_dims - 2; in HouseRow() local124 std::vector<int64> batch_dims(num_batch_dims); in HouseRow()125 for (int k = 0; k < num_batch_dims; ++k) { in HouseRow()189 const int64 num_batch_dims = num_dims - 2; in HouseCol() local190 std::vector<int64> batch_dims(num_batch_dims); in HouseCol()191 for (int k = 0; k < num_batch_dims; ++k) { in HouseCol()263 const int64 num_batch_dims = num_dims - 2; in HouseHolderBidiagonalization() local264 std::vector<int64> batch_dims(num_batch_dims); in HouseHolderBidiagonalization()265 for (int i = 0; i < num_batch_dims; ++i) { in HouseHolderBidiagonalization()468 const int64 num_batch_dims = num_dims - 2; in OneSidedJacobiUpdate() local[all …]
429 const int64 num_batch_dims = num_dims - 2; in SelfAdjointEig() local430 std::vector<int64> batch_dims(num_batch_dims); in SelfAdjointEig()431 for (int i = 0; i < num_batch_dims; ++i) { in SelfAdjointEig()
938 int64 num_batch_dims = in EmitBatchDotOperation() local947 CollapseFirstNDims(b, lhs_array, num_batch_dims); in EmitBatchDotOperation()949 CollapseFirstNDims(b, rhs_array, num_batch_dims); in EmitBatchDotOperation()951 CollapseFirstNDims(b, target_array, num_batch_dims); in EmitBatchDotOperation()975 dot_info.dim_nums.lhs_contracting_dimensions(0) - num_batch_dims); in EmitBatchDotOperation()978 dot_info.dim_nums.rhs_contracting_dimensions(0) - num_batch_dims); in EmitBatchDotOperation()