)
    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();
}