Method init

gemma/gm/nn/_transformer_like.py:88–115

Initializes a module method with variables and returns modified variables. ``init`` takes as first argument either a single ``PRNGKey``, or a dictionary mapping variable collections names to their ``PRNGKeys``, and will call ``method`` (which is the module's ``__call__`` function by

(
      self,
      rngs: PRNGKey | RNGSequences,
      *args,
      method: Callable[..., Any] | str | None = None,
      mutable: CollectionFilter = DenyList('intermediates'),
      capture_intermediates: (
          bool | Callable[[flax.linen.Module, str], bool]
      ) = False,
      **kwargs,
  )

  INFO: ClassVar[ModelInfo]

  def init(
      self,
      rngs: PRNGKey | RNGSequences,
      *args,
      method: Callable[..., Any] | str | None = None,
      mutable: CollectionFilter = DenyList('intermediates'),
      capture_intermediates: (
          bool | Callable[[flax.linen.Module, str], bool]
      ) = False,
      **kwargs,
  ) -> FrozenVariableDict | dict[str, Any]:
    """Initializes a module method with variables and returns modified variables.

    ``init`` takes as first argument either a single ``PRNGKey``, or a
    dictionary mapping variable collections names to their ``PRNGKeys``, and
    will call ``method`` (which is the module's ``__call__`` function by
    default) passing ``*args`` and ``**kwargs``, and returns
    a dictionary of initialized variables.

    Args:
      rngs: The PRNGKey or dictionary of PRNGKeys.
      *args: Positional arguments to pass to the method.
      method: The module method to initialize. Defaults to `__call__`.
      mutable: A filter for which variable collections are mutable.
      capture_intermediates: Whether to capture intermediate values.
      **kwargs: Keyword arguments to pass to the method.
    """
    ...

  @typechecked
  @abc.abstractmethod