// passing the platform index.
  FunctionType function_type = func_op.getFunctionType();
  ArrayRef<Type> new_input_types =
      function_type.getInputs().take_back(func_op.getNumArguments() - 1);
  func_op.setFunctionType(
      FunctionType::get(ctx, new_input_types, function_type.getResults()));
  func_op.getBody().eraseArgument(0);
}

}  // namespace

  }

  SmallVector<Type> output_types;
  for (const Value output : xla_call_module_op.getOutput()) {
    output_types.push_back(output.getType());
  }

  entry_func_op.setFunctionType(
      rewriter.getFunctionType(arg_types, output_types));

  // Replace argument types and locs.
  Block& entry = entry_func_op->getRegion(0).front();
  for (auto [arg, arg_type, arg_loc] :