Searched refs:lhs_non_contracting_dims (Results 1 – 7 of 7) sorted by relevance
57 dims.lhs_non_contracting_dims.push_back( in ParseConvolutionDimsInfo()89 dims.lhs_non_contracting_dims.push_back({lhs, rhs, output, i}); in ParseConvolutionDimsInfo()178 dnums.lhs_non_contracting_dims.emplace_back(); in ParseDotGeneralFromDot()179 dnums.lhs_non_contracting_dims.back().lhs = i; in ParseDotGeneralFromDot()180 dnums.lhs_non_contracting_dims.back().rhs = -1; in ParseDotGeneralFromDot()181 dnums.lhs_non_contracting_dims.back().output = in ParseDotGeneralFromDot()183 dnums.lhs_non_contracting_dims.size() - 1; in ParseDotGeneralFromDot()184 dnums.lhs_non_contracting_dims.back().spatial_dim = -1; in ParseDotGeneralFromDot()195 dnums.lhs_non_contracting_dims.size() + in ParseDotGeneralFromDot()
51 std::vector<int64> lhs_non_contracting_dims; in CanonicalizeDot() local52 lhs_non_contracting_dims.reserve(num_lhs_non_contracting_dims); in CanonicalizeDot()64 lhs_non_contracting_dims.push_back(i); in CanonicalizeDot()76 lhs_transpose.insert(lhs_transpose.end(), lhs_non_contracting_dims.begin(), in CanonicalizeDot()77 lhs_non_contracting_dims.end()); in CanonicalizeDot()
45 std::vector<DimNums> lhs_non_contracting_dims; member
355 : dnums.lhs_non_contracting_dims) { in InferDotShardingFromOperands()372 ? dnums.lhs_non_contracting_dims in InferDotShardingFromOperands()1054 : dnums.lhs_non_contracting_dims) { in InferDotOperandSharding()1068 for (const auto& dim : operand_index == 0 ? dnums.lhs_non_contracting_dims in InferDotOperandSharding()1084 for (const auto& dim : operand_index == 0 ? dnums.lhs_non_contracting_dims in InferDotOperandSharding()
63 mapping.lhs_non_contracting_dims.emplace_back(); in HandleDot()64 mapping.lhs_non_contracting_dims.back().lhs = i; in HandleDot()65 mapping.lhs_non_contracting_dims.back().rhs = -1; in HandleDot()66 mapping.lhs_non_contracting_dims.back().output = next_output_dim++; in HandleDot()369 for (const auto& mapping : dims_mapping.lhs_non_contracting_dims) { in ComputeDimensionIndexMapping()1833 std::vector<int64> lhs_non_contracting_dims; in PartitionDotGroupOnBatch() local1835 lhs_non_contracting_dims.reserve( in PartitionDotGroupOnBatch()1836 dims_mapping.lhs_non_contracting_dims.size()); in PartitionDotGroupOnBatch()1839 for (const auto& dim : dims_mapping.lhs_non_contracting_dims) { in PartitionDotGroupOnBatch()1840 lhs_non_contracting_dims.push_back(dim.lhs); in PartitionDotGroupOnBatch()[all …]
1004 for (const auto& dims : dims_info.lhs_non_contracting_dims) { in HandleConvolution()1005 mapping.lhs_non_contracting_dims.emplace_back(); in HandleConvolution()1006 mapping.lhs_non_contracting_dims.back().lhs = dims.lhs; in HandleConvolution()1007 mapping.lhs_non_contracting_dims.back().rhs = dims.rhs; in HandleConvolution()1008 mapping.lhs_non_contracting_dims.back().output = dims.output; in HandleConvolution()1009 mapping.lhs_non_contracting_dims.back().spatial = dims.spatial_dim; in HandleConvolution()
395 std::vector<DimsMapping> lhs_non_contracting_dims; member