class_docstr: Centralised, batched inference server.
class_docstr_tok: ["Centralised",",","batched","inference","server","."]
docstr: Replaces VariableReferences with their corresponding param values.
docstr_tok: ["Replaces","VariableReferences","with","their","corresponding","param","values","."]
code: def _dereference_params(self, arg):
      """Replaces VariableReferences with their corresponding param values."""
  
      if not isinstance(arg, variable_utils.VariableReference):
          # All arguments but VariableReference are returned without modifications.
          return arg
  
      # Due to batching dimension we take the first element.
      variable_name = arg.variable_name[0]
  
      if variable_name not in self._keys:
          # Create a new VariableClient which also serves new variables.
          self._keys.append(variable_name)
          self._variable_client = variable_utils.VariableClient(
              client=self._variable_source,
              key=self._keys,
              update_period=self._config.update_period,
          )
  
      if self._variable_client is None:
          raise ValueError("_variable_client not set")
  
      params = self._variable_client.params
      device_idx = self._call_cnt % len(self._devices)
      # Select device via round robin, and update its params if they changed.
      if self._device_params_ids[device_idx] != id(params):
          self._device_params_ids[device_idx] = id(params)
          self._device_params[device_idx] = jax.device_put(
              params, self._devices[device_idx]
          )
  
      # Return the params that are located on the chosen device.
      device_params = self._device_params[device_idx]
      if len(self._keys) == 1:
          return device_params
      return device_params[self._keys.index(variable_name)]
  
file_path: acme/acme/jax/inference_server.py
imports: import dataclasses
  import datetime
  import threading
  from typing import Any, Callable, Generic, Optional, Sequence, TypeVar
  import acme
  from acme.jax import variable_utils
  import jax
  import launchpad
label_desc: Use image node_id 3 for calling the InferenceServer obj's underlying member method code with example usage: obj._dereference_params(arg) and returns: device_params, arg, device_params