)
    if use_kernel:
      model.einsum = model.einsum_with_kernel
      model_inputs = {'x': x}
    else:
      model.einsum = model.einsum_without_kernel
      model_inputs = {'x': x, 'y': y}

    saved_model_save.save(
        model, self._input_saved_model_path, signatures=model.einsum
    )

    signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY

    MLIR pipeline in the future.
  }];
  let constructor = "TF::CreateGpuOpFusionPass()";
}

def BatchMatMulToEinsumPass : Pass<"tf-batch-matmul-to-tf-einsum", "mlir::func::FuncOp"> {
  let summary = "Replace TF BatchMatMul op by TF Einsum op.";
  let constructor = "TF::CreateBatchMatMulToEinsumPass()";
}

// EinsumOp
//===----------------------------------------------------------------------===//

// Verifies that,
// * Arity of the op is at most two.
//
// TODO(hinsu): Verify einsum equation attribute.
LogicalResult EinsumOp::verify() {
  EinsumOp op = *this;
  if (op.getN() > 2) {
    return op.emitOpError("supports at most two operands");
  }
  return success();
}

  func.return %0 : tensor<?x10xf32>
}

// -----

func.func @testTernaryEinsum(%arg0: tensor<2x3xf32>){
  // expected-error @+1 {{supports at most two operands}}
  %0 = "tf.Einsum"(%arg0, %arg0, %arg0) {equation = "ab,cd,ef->"} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> (tensor<*xf32>)
  func.return
}

// -----