#include "tensorflow/core/framework/resource_mgr.h"

namespace mlir {
namespace TF {

//===----------------------------------------------------------------------===//
// TensorFlow Contraction Fusion.
//===----------------------------------------------------------------------===//

struct ContractionFusion {
  explicit ContractionFusion(
      StringRef output_kernel, ArrayRef<int> additional_arguments = {},

        const int y_col =
            !matmul_op.getAdjY() ? shape_y.back() : *(shape_y.rbegin() + 1);

        // Checks that matrix multiply can perform a valid contraction.
        if (x_col != y_row) {
          result_shape.clear();
          return false;
        }

        result_shape.push_back(x_row);
        result_shape.push_back(y_col);
        return true;

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/util/matmul_bcast.h"

namespace mlir {
namespace TF {

// TF.Einsum provides fully general tensor contractions. For a few select
// cases, we can convert this op to other TF Ops, which in later passes
// properly convert to TF Lite ops.
struct ConvertTFEinsumOp : public OpRewritePattern<TF::EinsumOp> {
 public: