Searched defs:batch_dimensions (Results 1 – 5 of 5) sorted by relevance
57 const tensorflow::protobuf::RepeatedField<int64_t> &batch_dimensions = in HandleCustomCall() local
487 std::vector<int64_t> batch_dimensions; in BuildTriangularSolve() local
2092 HloInstruction* dot_operand, absl::Span<const int64_t> batch_dimensions, in NormalizeDotOperandToBatchMajorAndContractingMinor()
1618 DotDimensionsInfo(ShapedType type, ArrayRef<int64_t> batch_dimensions, in DotDimensionsInfo()1641 const DimensionVector &batch_dimensions() const { return batch_dimensions_; } in batch_dimensions() function in mlir::TF::__anonab43b57b0111::DotDimensionsInfo
444 auto batch_dimensions = in BatchDot() local3257 auto batch_dimensions = llvm::to_vector<4>(llvm::seq<int64_t>(0, rank - 2)); in matchAndRewrite() local