removed_params.push_back(false);
    }
  }
  // Remove old resource-type parameters.
  callee.getBody().front().eraseArguments(removed_params);
  // Update function type.
  callee.setFunctionType(FunctionType::get(callee.getContext(),
                                           callee.getBody().getArgumentTypes(),
                                           callee.getResultTypes()));
}

  // 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] :