arg.setType(new_arg_type);
    }

    // Update the function type.
    func.setType(mlir::FunctionType::get(module.getContext(),
                                         func.getBody().getArgumentTypes(),
                                         func.getFunctionType().getResults()));
  }
  return success();
}

}  // namespace tf_saved_model

                    esClient.admin().cluster().prepareNodesHotThreads().setIgnoreIdleThreads(ignoreIdleThreads).setInterval(interval)
                            .setThreads(threads).setTimeout(timeout).setType(type).execute().actionGet(timeout);
            append(buf, "cluster_name", () -> response.getClusterName().value()).append(',');
            final String hotThreads = response.getNodesMap().entrySet().stream().map(e -> {

		// This slice header is already defined in runtime/proc.go, so we update it here with new contents.
		sh := ldr.Lookup("runtime.runtime_inittasks", 0)
		sb := ldr.MakeSymbolUpdater(sh)
		sb.SetSize(0)
		sb.SetType(sym.SNOPTRDATA) // Could be SRODATA, but see issue 58857.
		sb.AddAddr(ctxt.Arch, t)
		sb.AddUint(ctxt.Arch, uint64(ldr.SymSize(t)/int64(ctxt.Arch.PtrSize)))
		sb.AddUint(ctxt.Arch, uint64(ldr.SymSize(t)/int64(ctxt.Arch.PtrSize)))

					Port:     svcInfo.NodePort(),
					Protocol: protocol,
					SetType:  utilipset.BitmapPort,
				}}
			case utilipset.ProtocolUDP:
				nodePortSet = proxier.ipsetList[kubeNodePortSetUDP]
				entries = []*utilipset.Entry{{
					// No need to provide ip info
					Port:     svcInfo.NodePort(),
					Protocol: protocol,
					SetType:  utilipset.BitmapPort,
				}}
			case utilipset.ProtocolSCTP:

      Value handle = tensor_array_v3_op.getHandle();
      auto handle_type = mlir::cast<TensorType>(handle.getType());
      if (handle_type.getElementType().isInteger(/*width=*/32)) {
        Type resource_tensor_type =
            handle_type.clone(TF::ResourceType::get(rewriter.getContext()));
        handle.setType(resource_tensor_type);
      }
    }

    // Special type fixes for scalar tensor types.

    TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));

    auto metadata_arg = metadata_proto.add_args();
    metadata_arg->set_kind(tpu::TPUCompileMetadataProto::Arg::PARAMETER);
    metadata_arg->set_dtype(dtype);
  }

  return absl::OkStatus();
}

absl::Status SetupReturnValues(mlir::ModuleOp module,
                               tpu::TPUCompileMetadataProto& metadata_proto) {

    int64_t index = entry.getFirst();
    Type new_type = entry.getSecond().first;
    int64_t updated_index = entry.getSecond().second;
    auto operand = caller->getOperand(index);
    builder.setInsertionPoint(caller);
    new_operands[index] = builder.create<TF::ReadVariableOp>(
        caller->getLoc(), ArrayRef<Type>{new_type}, ArrayRef<Value>{operand});
    caller->setOperand(index, new_operands[index]);

    new_quantized_func.setType(
        FunctionType::get(getContext(), TypeRange{ValueRange{args}},
                          new_quantized_func.getResultTypes()));
    for (auto [partitioned_call_arg, new_quantized_func_arg] :
         llvm::zip_equal(args, new_quantized_func.getArguments())) {
      new_quantized_func_arg.setType(partitioned_call_arg.getType());
    }

    SmallVector<Value> new_float_func_args{func_args.begin(), func_args.end()};
    new_float_func_args[weight_operand_idx] = reshape_op;
    new_float_func.getArgument(weight_operand_idx).setType(new_shape);
    new_float_func.setType(FunctionType::get(
        getContext(), TypeRange{ValueRange{new_float_func_args}},
        new_float_func.getResultTypes()));
    symbol_table.insert(new_float_func);

          block.getArgument(it.first->getSecond()));
    }
    var_handle_op.erase();
  }

  if (!var_handle_shared_names->empty())
    function.setType(FunctionType::get(function.getContext(), func_arg_types,
                                       func_type.getResults()));

  return success();
}

// Records the current live value for a resource variable and whether a read or