if (!main_fn) {
    return absl::InternalError("Could not find main function in MLIR Module.");
  }

  mlir::FunctionType func_type = main_fn.getFunctionType();
  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;
    xla::Shape xla_shape = xla::TypeToShape(input_type);
    TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(
        xla_shape.dimensions(), &tensor_shape));

                            ConversionPatternRewriter *rewriter) const {
    auto guard = OpBuilder::InsertionGuard(*rewriter);
    auto inputs = branch_func.getFunctionType().getInputs();
    Block *block = rewriter->createBlock(
        &branch_func.getBody(), branch_func.begin(), inputs,
        SmallVector<Location>(inputs.size(), branch_func.getLoc()));

    if (op.getN() == 1) {
      rewriter.replaceOpWithNewOp<UnaryEinsumOp>(
          op, op.getType(), *op.getInputs().begin(), equation);
    } else if (op.getN() == 2) {
      ValueRange inputs = op.getInputs();
      rewriter.replaceOpWithNewOp<EinsumOp>(op, op.getType(), inputs[0],
                                            inputs[1], equation);
    } else {

          while_op.getInput().getTypes(), while_op.getOutput().getTypes(),
          while_op.ResolveBodyFunction(&symbol_table_)
              .getFunctionType()
              .getInputs());
      return PropagateShapeToFunctions(
          module, compatible_types,
          {while_op.ResolveCondFunction(&symbol_table_),
           while_op.ResolveBodyFunction(&symbol_table_)},