OpBuilder builder(func.getBody());
  std::string empty_option_buffer;
  auto op = builder.create<CustomOp>(
      func.getLoc(), func.getFunctionType().getResults(), func.getArguments(),
      api, CustomOption(&builder, empty_option_buffer));
  builder.create<func::ReturnOp>(func.getLoc(), op.getResults());
  return success();
}

LogicalResult VerifyNgrams(func::FuncOp func) {

  assert(caller_region.empty() &&
         "Expected empty region for newly created ops");
  OpBuilder builder(caller_region);
  Block* entry = builder.createBlock(&caller_region);

  auto loc = op->getLoc();
  if (use_region_args) {
    auto inputs = func.getFunctionType().getInputs();
    entry->addArguments(inputs, SmallVector<Location>(inputs.size(), loc));
    args = entry->getArguments();
  }

  auto reshape_sizes = DenseIntElementsAttr::get(reshape_type, new_shape);
  auto reshape_value =
      builder->create<TF::ConstOp>(input.getLoc(), reshape_sizes);
  return builder->create<TF::ReshapeOp>(input.getLoc(), reshape_result_type,
                                        input, reshape_value);
}

// Creates transpose op for shape to depth transform.

        if (IsCompatibleTypeWithTFLCastOp(value.getType()) &&
            IsCompatibleTypeWithTFLCastOp(type)) {
          auto cast = b.create<CastOp>(yield_op->getLoc(), type, value);
          args.push_back(cast);
        } else {
          auto cast = b.create<TF::CastOp>(yield_op->getLoc(), type, value);
          args.push_back(cast);
        }
      }
    }
    args.append(new_args.begin(), new_args.end());
  } else {

    auto var_handle_op =
        builder.create<TF::VarHandleOp>(const_op.getLoc(),
                                        /*resource=*/resource_type,
                                        /*container=*/"", shared_name);

    auto read_variable_op = builder.create<TF::ReadVariableOp>(
        const_op.getLoc(), const_op.getType(), var_handle_op);

        mlir::cast<TensorType>(input.getType()), kNchwToNhwcPermutation);

    auto input_transpose_op = rewriter.create<mlir::stablehlo::TransposeOp>(
        op.getLoc(), /*resultType0=*/new_input_tensor_type, /*operand=*/input,
        rewriter.getDenseI64ArrayAttr(kNchwToNhwcPermutation));

    // Transpose the filter tensor: [o, i, 0, 1] => [0, 1, i, o]
    Value filter = op->getOperand(1);

  }

  auto compile_op = builder->create<TF::_TPUCompileMlirOp>(
      cluster_func.getLoc(),
      /*compilation_status=*/compilation_status_type, /*program=*/
      llvm::SmallVector<Type, 8>(num_cores_per_replica, program_type),
      compile_op_operands, txt_module, txt_metadata);

  return tensorflow::WrapOpInLaunch(builder, compile_op.getLoc(), compile_op,
                                    compilation_device);
}

  // Unformat op.
  auto unformat_op = builder.create<TF::TPUReshardVariablesOp>(
      while_op.getLoc(), llvm::ArrayRef<Type>{}, unformat_operands);
  WrapOpInLaunch(&builder, execute_launch.getLoc(), unformat_op,
                 execute_launch.getDevice());
  builder.create<tf_device::ReturnOp>(while_op.getLoc(), ArrayRef<Value>{});

  return true;
}

void TPUVariableRuntimeReformattingPass::runOnOperation() {

      block_arg.getLoc(),
      DenseIntElementsAttr::get(
          RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
          group_assignment_val));

  StringAttr reduce_op = builder.getStringAttr("Add");
  StringAttr mode = builder.getStringAttr("CrossReplica");
  return builder.create<XlaAllReduceOp>(block_arg.getLoc(), block_arg.getType(),

    auto *new_ewise_op =
        rewriter.create(op->getLoc(), op->getName().getIdentifier(),
                        {tpose_arg1.getOperand(0), tpose_arg2.getOperand(0)},
                        new_out_type, op->getAttrs());

    // Apply original tranpose to output of ewise op.
    auto out_tpose_op = rewriter.create<TFL::TransposeOp>(
        new_ewise_op->getLoc(), op->getResult(0).getType(),