Method apply

gemma/gm/nn/_transformer_like.py:134–145

Applies a module method to variables and returns output and modified variables.

(
      self,
      variables: VariableDict,
      tokens: Int['*B L'],
      *,
      images: UInt8['*B N H W C'] | UInt8['*B H W C'] | None = None,
      cache: _config.Cache | None = None,
      positions: Int['*B L_with_mm'] | None = None,
      attention_mask: Bool['*B L_with_mm cache_length'] | None = None,
  )

  @typechecked
  @abc.abstractmethod
  def apply(
      self,
      variables: VariableDict,
      tokens: Int['*B L'],
      *,
      images: UInt8['*B N H W C'] | UInt8['*B H W C'] | None = None,
      cache: _config.Cache | None = None,
      positions: Int['*B L_with_mm'] | None = None,
      attention_mask: Bool['*B L_with_mm cache_length'] | None = None,
  ) -> Any | tuple[Any, FrozenVariableDict | dict[str, Any]]:
    """Applies a module method to variables and returns output and modified variables."""
    ...

  @typechecked
  @abc.abstractmethod