# Copyright 2024 The Flax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Flax functional core: Scopes."""

import collections
import contextlib
import dataclasses
import functools
import hashlib
import typing
from typing import (
  Any,
  Callable,
  Dict,
  Generic,
  Iterable,
  Literal,
  Mapping,
  Optional,
  Sequence,
  Set,
  Tuple,
  TypeVar,
  Union,
  cast,
  overload,
)

import jax
import numpy as np
from jax import numpy as jnp
from jax import random, tree_util

from flax import config as config
from flax import configurations as legacy_config  # only for flax_lazy_rng
from flax import errors, struct, traceback_util
from flax.ids import uuid
from flax.typing import (
  PRNGKey,
  Array,
  RNGSequences,
  Collection,
  MutableCollection,
  VariableDict,
  FrozenVariableDict as FrozenVariableDict,
  MutableVariableDict,
  PRNGFoldable,
)

from . import meta, partial_eval, tracers
from .frozen_dict import FrozenDict, freeze, unfreeze

traceback_util.register_exclusion(__file__)

T = TypeVar('T')


Filter = Union[bool, str, typing.Collection[str], 'DenyList']

# When conditioning on filters we require explicit boolean comparisons.
# pylint: disable=g-bool-id-comparison


@dataclasses.dataclass(frozen=True, eq=True)
class DenyList:
  """DenyList represents an opt-out based mutability filter.
  DenyList can be used to make every collection mutable except the ones
  defined in the given filter.
  To for example make everything but the params collection mutable::
    nn.apply(fn, mutable=nn.DenyList(["params"]))
  Attributes:
    deny: The filter representing the collections that are not mutable.
  """

  deny: Filter


CollectionFilter = Filter
PRNGSequenceFilter = Filter


class LazyRng(struct.PyTreeNode):
  """Wrapper around JAX PRNGKey that lazily maintains a tuple of static data to be folded into the rng."""

  rng: PRNGKey
  suffix: Tuple[PRNGFoldable, ...] = struct.field(pytree_node=False)

  def as_jax_rng(self) -> PRNGKey:
    return _fold_in_static(self.rng, self.suffix)

  @staticmethod
  def create(
    rng: Union['LazyRng', PRNGKey], *suffix: PRNGFoldable
  ) -> 'LazyRng':
    if not legacy_config.flax_lazy_rng:
      if isinstance(rng, LazyRng):
        assert not rng.suffix
        rng = rng.rng
      return LazyRng(_legacy_rng_fold_in(rng, suffix), ())
    if isinstance(rng, LazyRng):
      return LazyRng(rng.rng, rng.suffix + suffix)
    else:
      return LazyRng(rng, suffix)


def _legacy_rng_fold_in(rng: PRNGKey, data: Iterable[PRNGFoldable]) -> PRNGKey:
  """Legacy RNG folding."""
  for x in data:
    if isinstance(x, str):
      m = hashlib.sha1()
      m.update(x.encode('utf-8'))
      d = m.digest()
      hash_int = int.from_bytes(d[:4], byteorder='big')
      rng = random.fold_in(rng, jnp.uint32(hash_int))  # type: ignore
    elif isinstance(x, int):
      rng = random.fold_in(rng, x)
    else:
      raise ValueError(f'Expected int or string, got: {x}')
  return rng


def _fold_in_static(
  rng: PRNGKey, data: typing.Collection[PRNGFoldable]
) -> PRNGKey:
  """Folds static data (strings & ints) into a jax.random.PRNGKey using its SHA-1 hash.

  This is faster than splitting an PRNGKey because it allows generating new PRNG
  keys in parallel that are independent of each other.

  Args:
   rng: the rng to fold the string into.
   data: the string to be folded in.

  Returns:
   The newly generated PRNG key.
  """
  if not data:
    return rng
  m = hashlib.sha1()
  for x in data:
    if config.flax_fix_rng_separator:
      # encode seperate to avoid collisions like for example: ("ab", "c") and ("a", "bc")
      m.update(b'\00')
    if isinstance(x, str):
      m.update(x.encode('utf-8'))
    elif isinstance(x, int):
      m.update(x.to_bytes((x.bit_length() + 7) // 8, byteorder='big'))
    else:
      raise ValueError(f'Expected int or string, got: {x}')
  d = m.digest()
  hash_int = int.from_bytes(d[:4], byteorder='big')
  return random.fold_in(rng, jnp.uint32(hash_int))  # type: ignore


def is_filter_empty(filter_like: Filter) -> bool:
  """Returns True if `filter_like` is an empty filter.

  Args:
    filter_like: The filter to test.

  Returns:
    A filter is empty when it is an empty collection, it is a bool with value
    False, ir it is a DenyList that matches everything. A string filter is never
    empty.
  """
  if isinstance(filter_like, str):
    return False
  if isinstance(filter_like, typing.Collection):
    return not filter_like
  if isinstance(filter_like, bool):
    return not filter_like
  if isinstance(filter_like, DenyList):
    # if any arbitrary collection is in the denylist it matches everything so
    # the filter is empty. This is checked with a stub.
    return in_filter(filter_like.deny, '__flax_internal_stub__')
  raise errors.InvalidFilterError(filter_like)


def in_filter(filter_like: Filter, col: str) -> bool:
  """Checks whether a filter can be applied to a collection.

  Used for both collections and rng sequence filters.

  Args:
    filter_like: a filter (either a boolean, a string, or a list of strings) for
      a collection.
    col: a collection, which is a string identifying a dictionary of data, for
      instance "params" or "batch_stats".

  Returns:
    True if either `filter_like` is True, equal to `col`, or a sequence
    containing `col`.
  """
  if isinstance(filter_like, str):
    return col == filter_like
  if isinstance(filter_like, typing.Collection):
    return col in filter_like
  if isinstance(filter_like, bool):
    return filter_like
  if isinstance(filter_like, DenyList):
    return not in_filter(filter_like.deny, col)
  raise errors.InvalidFilterError(filter_like)


def filter_to_set(x: Filter) -> Set[str]:
  """Converts a Filter into a set of collections, fails on the infinite set.

  Args:
    x: a filter (boolean, string, or list of strings).

  Returns:
    The input filter represented as a set of strings.
  """
  assert x is not True and not isinstance(x, DenyList), 'Infinite set'
  if x is False:
    return set()
  if isinstance(x, str):
    return set([x])
  if isinstance(x, typing.Collection):
    return set(x)
  raise errors.InvalidFilterError(x)


def union_filters(a: Filter, b: Filter) -> Filter:
  """Takes the union of two filters (similar to a logical or).

  Args:
    a: a filter.
    b: a filter.

  Returns:
    The union of the two input filters. For instance,
    `union_filters('f1', ['f2']) = {'f1', 'f2'}`.
  """
  if a is True or b is True:
    return True
  if isinstance(a, DenyList) and isinstance(b, DenyList):
    return DenyList(intersect_filters(a.deny, b.deny))
  if isinstance(b, DenyList):
    a, b = b, a
  if isinstance(a, DenyList):
    return DenyList(subtract_filters(a.deny, b))

  a = filter_to_set(a)
  b = filter_to_set(b)
  return a.union(b)


def subtract_filters(a: Filter, b: Filter) -> Filter:
  """Returns the subtraction of b from a.

  Args:
    a: a filter.
    b: a filter.

  Returns:
    A filter matching with values in a that are not in b.
  """
  if b is True:
    return False
  if a is True:
    return DenyList(b)
  if isinstance(a, DenyList) and isinstance(b, DenyList):
    return subtract_filters(b.deny, a.deny)
  if isinstance(a, DenyList):
    return DenyList(union_filters(a.deny, b))
  if isinstance(b, DenyList):
    return intersect_filters(a, b.deny)
  a = filter_to_set(a)
  b = filter_to_set(b)
  return a - b


def intersect_filters(a: Filter, b: Filter) -> Filter:
  """Take the intersection of two filters (similar to a logical and).

  Args:
    a: a filter.
    b: a filter.

  Returns:
    The intersection of the two input filters. For instance,
    `intersect_filters('f1', ['f1', 'f2']) = {'f1'}`.
  """
  if a is True:
    return b
  if b is True:
    return a
  if isinstance(a, DenyList) and isinstance(b, DenyList):
    return DenyList(union_filters(b.deny, a.deny))
  if isinstance(b, DenyList):
    b, a = a, b
  if isinstance(a, DenyList):
    return subtract_filters(b, a.deny)
  a = filter_to_set(a)
  b = filter_to_set(b)
  return a.intersection(b)


def group_collections(
  xs: VariableDict, col_filters: Sequence[CollectionFilter]
) -> Sequence[MutableVariableDict]:
  """Groups variables by collection filters.

  Iteratively applies the filters in `col_filters` to `xs`, and adds the result
  of applying each filter to the output sequence. Each key in `xs` is only added
  to the output once.

  Args:
    xs: a dictionary of variables, keyed by collections (strings).
    col_filters: a list of collection filters.

  Returns:
    A sequence S with `len(S) == len(col_filters)`. Each `S[i]` is the result of
    applying filter `col_filters[i]` to the remaining keys in `xs`.
  """
  cols: Iterable[str]
  cols = xs.keys()
  groups = []
  for col_filter in col_filters:
    remaining_cols = []
    group = {}
    for col in cols:
      if in_filter(col_filter, col):
        group[col] = jax.tree_util.tree_map(lambda x: x, xs[col])
      else:
        remaining_cols.append(col)
    cols = remaining_cols
    groups.append(group)
  return tuple(groups)


class Variable(Generic[T]):
  """A Variable object allows mutable access to a variable in a VariableDict.

  Variables are identified by a collection (e.g., "batch_stats") and a name
  (e.g., "moving_mean"). The value property gives access to the variable's
  content and can be assigned to for mutation.
  """

  def __init__(self, scope: 'Scope', collection: str, name: str, unbox: bool):
    """Initializes a variable.

    Args:
      scope: The scope in which the variable is stored.
      collection: The collection of the variable (e.g., "params").
      name: The name of the variable (e.g., "dense").
      unbox: Whether to unbox boxed values with metadata.
    """
    self._id = uuid()
    self.scope = scope
    self.collection = collection
    self.name = name
    self.unbox = unbox

  @property
  def value(self) -> T:
    """Returns the value of this Variable."""
    v = self.scope.get_variable(self.collection, self.name)
    return meta.unbox(v) if self.unbox else v

  @value.setter
  def value(self, value: T):
    """Updates the value of this Variable."""
    if self.unbox:
      cur = self.scope.get_variable(self.collection, self.name)
      cur_struct = tree_util.tree_structure(cur, is_leaf=meta.is_axis_metadata)
      value_struct = tree_util.tree_structure(
        value, is_leaf=meta.is_axis_metadata
      )
      has_meta = any(map(meta.is_axis_metadata, cur_struct.flatten_up_to(cur)))
      if cur_struct == value_struct and has_meta:
        value = meta.replace_boxed(cur, value)

    self.scope.put_variable(self.collection, self.name, value)

  def is_mutable(self) -> bool:
    """Checks if this Variable is mutable."""
    return self.scope.is_mutable_collection(self.collection)


class _ChildRNGSentinel:
  pass


# used to identify that an rng counter is meant for a child scope
child_rng_token = _ChildRNGSentinel()


class _DefaultSentinel:
  pass


# used to denote no default flag value on scope
no_flag = _DefaultSentinel()


class Scope:
  """A Scope allows easy access to variables and manages RNGS of a neural network layer.

  Scopes are purely functional and encapsulated in
  :class:`flax.linen.module.Module`, so users writing neural network code
  usually generally do not interact with ``Scopes`` directly.

  See `core design tests
  <https://github.com/google/flax/tree/main/tests/core/design>`_
  for a number of examples using ``Scopes``.
  """

  reservations: Dict[str, Set[Optional[str]]]

  def __init__(
    self,
    variables: MutableVariableDict,
    rngs: Optional[Union[RNGSequences, Dict[str, LazyRng]]] = None,
    name: Optional[str] = None,
    mutable: CollectionFilter = False,
    parent: Optional['Scope'] = None,
    path: Iterable[str] = (),
    debug_path: Iterable[str] = (),
    flags: Optional[Mapping] = None,
  ):
    """Initializes a Scope.

    Args:
      variables: VariableDict to initialize the Scope with.
      rngs: RNGs used in this scope or one of the child scopes.
      name: name of this scope.
      mutable: A CollectionFilter determining which variables are mutable.
      parent: The parent scope.
      path: The path in the variable tree from the root scope to this scope. It
        exactly matches the module path.
      debug_path: Similar to path but could contain transformation decorators.
      flags: internal flags.
    """
    rngs = {k: LazyRng.create(v) for k, v in rngs.items()} if rngs else {}
    self._variables = variables
    self.parent = parent
    self.name = name
    self.path = tuple(path)
    self.debug_path = tuple(debug_path) or self.path
    self.rngs = rngs
    self.mutable = mutable
    self.flags = freeze({} if flags is None else flags)

    self._root = parent.root if parent else None
    self.trace_level = tracers.trace_level(tracers.current_trace())

    self.rng_counters = {key: 0 for key in self.rngs}
    self.reservations = collections.defaultdict(set)

    self._invalid = False

  def __eq__(self, other: Any) -> bool:
    # If the root variable dict and path are the same, then two scopes behave
    # identically. Effectively, a scope is nothing more than a cursor into a
    # variable dict and an rng counter dict.
    if not isinstance(other, Scope):
      return False
    if self is other:
      return True
    return (
      self.root._variables is other.root._variables
      and self.path == other.path
      and self.rng_counters is other.rng_counters
    )

  def __hash__(self) -> int:
    # see __eq__
    return hash((id(self.root._variables), self.path, id(self.rng_counters)))

  @property
  def root(self) -> 'Scope':
    return self._root or self

  @property
  def path_text(self) -> str:
    """Returns the debug path as a human readable string."""
    return '/' + '/'.join(self.debug_path)

  @property
  def invalid(self) -> bool:
    """Returns true if this scope is invalidated as a result of `Scope.temporary`."""
    return self._invalid

  def _check_valid(self):
    if self._invalid:
      raise errors.InvalidScopeError(self.name)

  @contextlib.contextmanager
  def temporary(self):
    """Returns a context manager that will invalidate this Scope when leaving the context."""
    try:
      yield self
    finally:
      self.invalidate()

  def invalidate(self):
    """Invalidates the Scope."""
    self._invalid = True

  def mutable_variables(self) -> Union[VariableDict, Dict[str, Any]]:
    """Returns an immutable copy of the mutable variables belonging to this Scope."""
    self._populate_collections()
    xs = {
      k: v for k, v in self._variables.items() if in_filter(self.mutable, k)
    }
    if config.flax_return_frozendict:
      return freeze(xs)
    return xs

  def variables(self) -> Union[VariableDict, Dict[str, Any]]:
    """Returns an immutable copy of the variables belonging to this Scope."""
    self._populate_collections()
    if config.flax_return_frozendict:
      return freeze(self._variables)
    return self._variables

  def _validate_trace_level(self):
    tracers.check_trace_level(self.trace_level)

  def rewound(self, rewind_rngs: bool = False) -> 'Scope':
    """Returns a rewound version of this Scope.

    Args:
      rewind_rngs: if true, reset the RNG counter of this scope.

    Returns:
      A rewound version of this scope, which means reservations are
      emptied, and the rng counter is optionally rewound.
    """
    self._check_valid()
    scope = Scope(
      self._variables,
      self.rngs,
      self.name,
      self.mutable,
      self.parent,
      path=self.path,
      debug_path=self.debug_path,
      flags=self.flags,
    )
    if not rewind_rngs:
      scope.rng_counters = self.rng_counters
    return scope

  def name_reserved(self, name: str, col: Optional[str] = None) -> bool:
    """Checks whether a name for a child Scope or Variable is taken.

    Args:
      name: the name to check for collision.
      col: if a variable, the collection used.
    """
    if name in self.reservations:
      # allow the same name for two variables in
      # different collections, otherwise raise error.
      if (
        None in self.reservations[name]
        or col is None
        or col in self.reservations[name]
      ):
        return True
    return False

  def reserve(self, name: str, col: Optional[str] = None):
    """Reserves a name for a child Scope or Variable.

    Throws an error if the name exists already.

    Args:
      name: the name to reserve.
      col: if a variable, the collection used.
    """
    if not isinstance(name, str):
      raise TypeError(
        'The type of scope "{name}" should be string but ' f'it is {type(name)}'
      )
    if self.name_reserved(name, col):
      raise ValueError(f'Duplicate use of scope name: "{name}"')
    self.reservations[name].add(col)

  def default_name(self, prefix: str) -> str:
    """Generates an unreserved name with the given prefix.

    Args:
      prefix: prefix to use for generating an unreserved name.

    Returns:
      The generated name.
    """
    i = 0
    while True:
      name = f'{prefix}{i}'
      if name not in self.reservations:
        return name
      i += 1

  def push(
    self, name: Optional[str] = None, prefix: str = '', reuse=False
  ) -> 'Scope':
    """Creates a child Scope.

    Args:
      name: optional name of the child.
      prefix: prefix used for generating the name if `name` is `None`.
      reuse: if True will return a pre-existing child scope with the given name
        instead of throwing an error.

    Returns:
      The child scope.
    """
    self._check_valid()
    self._validate_trace_level()
    if name is None:
      name = self.default_name(prefix)
    if not reuse or name not in self.reservations:
      self.reserve(name)
    rngs = {key: LazyRng.create(rng, name) for key, rng in self.rngs.items()}
    rng_key = (child_rng_token, name)
    if rng_key in self.rng_counters:
      rng_counters = self.rng_counters.get(rng_key)  # type: ignore
    else:
      rng_counters = {key: 0 for key in rngs}
      self.rng_counters[rng_key] = rng_counters  # type: ignore
    scope = Scope(
      {},
      name=name,
      rngs=rngs,
      parent=self,
      mutable=self.mutable,
      path=self.path + (name,),
      debug_path=self.debug_path + (name,),
      flags=self.flags,
    )
    scope.rng_counters = rng_counters
    return scope

  def child(
    self,
    fn: Callable[..., Any],
    name: Optional[str] = None,
    prefix: Optional[str] = None,
    named_call: bool = True,
    **partial_kwargs,
  ) -> Callable[..., Any]:
    """Partially applies a child scope to fn.

    When calling the returned function multiple times variables will be reused.

    Args:
      fn: the function to partially apply the child Scope to.
      name: optional name of the child.
      prefix: prefix used for generating name if it is `None`.
      named_call: if true, `fn` will be run under `jax.named_scope`. The XLA
        profiler will use this to name tag the computation.
      **partial_kwargs: additional kwargs partially applied to `fn`.

    Returns:
      The function with a partially applied scope.
    """
    if name is None:
      if prefix is None:
        prefix = fn.__name__ + '_' if hasattr(fn, '__name__') else ''
      name = self.default_name(prefix)
    scope = self.push(name)

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
      kwargs = dict(partial_kwargs, **kwargs)
      if named_call:
        with jax.named_scope(name):
          res = fn(scope.rewound(), *args, **kwargs)
      else:
        res = fn(scope.rewound(), *args, **kwargs)
      return res

    return wrapper

  def is_mutable_collection(self, col: str) -> bool:
    """Returns true if the collection `col` is mutable."""
    return in_filter(self.mutable, col)

  def is_collection_empty(self, col: str) -> bool:
    """Returns true if the collection is empty."""
    if col in self.root._variables:  # pylint: disable=protected-access
      return not self.root._variables[col]  # pylint: disable=protected-access
    return True

  def _mutable_collection(self, col: str) -> MutableCollection:
    """Returns the collection `col` as a mutable object."""
    assert self.is_mutable_collection(col), f'Collection {col} is not mutable'

    # The actual variable dict is stored in the root scope only, and subscopes
    # hold references to subtrees relevant to them. This function ensures that
    # the collections are created in the top-level Scope and we return the
    # correct reference.
    if col not in self._variables:
      if not self.parent:
        # If this is the top-level Scope, just add an empty collection.
        self._variables[col] = {}
      else:
        assert self.name is not None  # Only top-level Scope have name None.
        # Populate the parent collections recursively and obtain a reference to
        # the direct parent (which, by transitivity, is be a reference to a
        # dict in the root Scope).
        parent_col = self.parent._mutable_collection(col)  # pylint: disable=protected-access
        if self.name not in parent_col:
          # If this Scope's name does not occur in the parent collection, add it
          # to the parent scope (updating the parent's variable dict).
          parent_col[self.name] = {}
        # Store a reference to the parent's scope collection for in this scope's
        # variable dict.
        self._variables[col] = parent_col[self.name]

    return self._variables[col]

  def _collection(self, col: str) -> Collection:
    """Returns a collection of variables of collection `col`."""
    if col not in self._variables:
      if self.parent:
        assert self.name is not None
        parent_col = self.parent._collection(col)  # pylint: disable=protected-access
        if self.name not in parent_col:
          return FrozenDict()
        self._variables[col] = parent_col[self.name]
      else:
        return FrozenDict()
    return self._variables[col]

  def has_rng(self, name: str) -> bool:
    """Returns true if a PRNGSequence with name `name` exists."""
    return name in self.rngs

  def make_rng(self, name: str = 'params') -> PRNGKey:
    """Generates A PRNGKey from a PRNGSequence with name `name`."""
    if not self.has_rng(name):
      if self.has_rng('params'):
        name = 'params'
      else:
        raise errors.InvalidRngError(f'{self.name} needs PRNG for "{name}"')
    self._check_valid()
    self._validate_trace_level()
    self.rng_counters[name] += 1
    return LazyRng.create(self.rngs[name], self.rng_counters[name]).as_jax_rng()

  def get_variable(self, col: str, name: str, default: Any = None) -> Any:
    """Retrieves the value of a Variable.

    Args:
      col: the variable collection.
      name: the name of the variable.
      default: the default value to return if the variable does not exist in
        this scope.

    Returns:
      The value of the input variable, of the default value if the variable
      doesn't exist in this scope.
    """
    variables = self._collection(col)
    if name in variables:
      return variables[name]
    else:
      return default

  def has_variable(self, col: str, name: str) -> bool:
    """Returns true if the given variable exists in this scope.

    Args:
      col: the collection of the variable.
      name: the name of the variable.
    """
    variables = self._collection(col)
    return name in variables

  def put_variable(self, col: str, name: str, value: Any):
    """Updates the value of the given variable if it is mutable, or an error otherwise.

    Args:
      col: the collection of the variable.
      name: the name of the variable.
      value: the new value of the given variable.
    """
    self._check_valid()
    self._validate_trace_level()
    if not self.is_mutable_collection(col):
      raise errors.ModifyScopeVariableError(col, name, self.path_text)
    variables = self._mutable_collection(col)

    # Make sure reference sharing of child variable dictionaries isn't broken.
    # See https://github.com/google/flax/issues/2022 for more details.
    def put(target, key, val):
      if (
        key in target
        and isinstance(target[key], dict)
        and isinstance(val, Mapping)
      ):
        for k, v in val.items():
          put(target[key], k, v)
      else:
        target[key] = val

    put(variables, name, value)

  @overload
  def variable(
    self,
    col: str,
    name: str,
    init_fn: Optional[Callable[..., T]] = None,
    *init_args,
  ) -> Variable[T]:
    ...

  @overload
  def variable(
    self,
    col: str,
    name: str,
    init_fn: Optional[Callable[..., T]] = None,
    *init_args,
    unbox: Literal[True],
    **init_kwargs,
  ) -> Variable[T]:
    ...

  @overload
  def variable(
    self,
    col: str,
    name: str,
    init_fn: Optional[Callable[..., T]] = None,
    *init_args,
    unbox: Literal[False],
    **init_kwargs,
  ) -> Variable[meta.AxisMetadata[T]]:
    ...

  @overload
  def variable(
    self,
    col: str,
    name: str,
    init_fn: Optional[Callable[..., T]] = None,
    *init_args,
    unbox: bool = True,
    **init_kwargs,
  ) -> Union[Variable[T], Variable[meta.AxisMetadata[T]]]:
    ...

  def variable(
    self,
    col: str,
    name: str,  # pylint: disable=keyword-arg-before-vararg
    init_fn: Optional[Callable[..., T]] = None,
    *init_args,
    unbox: bool = True,
    **init_kwargs,
  ) -> Union[Variable[T], Variable[meta.AxisMetadata[T]]]:
    """Creates a variable if it doesn't exist yet in this scope and returns it.

    Args:
      col: the collection of the variable.
      name: the name of the variable.
      init_fn: a function taking a PRNGKey plus any other number of positional
        arguments. If None, the variable must already be initialized otherwise
        an error is raised.
      *init_args: the positional arguments to evaluate init_fn on lazily.
      unbox: If True, ``AxisMetadata`` instances are replaced by their unboxed
        value, see ``flax.nn.meta.unbox`` (default: True).
      **init_kwargs: the key-word arguments to evaluate init_fn on lazily.

    Returns:
      The variable.  Throws an error if the variable exists already.
    """
    self.reserve(name, col)
    if not self.has_variable(col, name):
      if not self.is_mutable_collection(col) or init_fn is None:
        if self.is_collection_empty(col):
          raise errors.ScopeCollectionNotFound(col, name, self.path_text)
        raise errors.ScopeVariableNotFoundError(name, col, self.path_text)
      init_value = init_fn(*init_args, **init_kwargs)
      self.put_variable(col, name, init_value)
    # cast to make static analyzers happy
    return cast(
      Union[Variable[T], Variable[meta.AxisMetadata[T]]],
      Variable(self, col, name, unbox=unbox),
    )

  @overload
  def param(
    self, name: str, init_fn: Callable[..., T], *init_args,
  ) -> T:
    ...

  @overload
  def param(
    self,
    name: str,
    init_fn: Callable[..., T],
    *init_args,
    unbox: Literal[True],
    **init_kwargs,
  ) -> T:
    ...

  @overload
  def param(
    self,
    name: str,
    init_fn: Callable[..., T],
    *init_args,
    unbox: Literal[False],
    **init_kwargs,
  ) -> meta.AxisMetadata[T]:
    ...

  @overload
  def param(
    self,
    name: str,
    init_fn: Callable[..., T],
    *init_args,
    unbox: bool,
    **init_kwargs,
  ) -> Union[T, meta.AxisMetadata[T]]:
    ...

  def param(
    self,
    name: str,
    init_fn: Callable[..., T],
    *init_args,
    unbox: bool = True,
    **init_kwargs,
  ) -> Union[T, meta.AxisMetadata[T]]:
    """Creates a parameter if it doesn't exist yet in this scope and returns it.

    If the parameter exists already, the existing value is simply returned.

    Args:
      name: the name of the parameter.
      init_fn: a function taking a PRNGKey plus any other number of positional
        arguments.
      *init_args: the positional arguments to evaluate init_fn on lazily.
      unbox: If True, ``AxisMetadata`` instances are replaced by their unboxed
        value, see ``flax.nn.meta.unbox`` (default: True).
      **init_kwargs: the key-word arguments to evaluate init_fn on lazily.

    Returns:
      The parameters. Throws an error if the params exist already.
    """
    self.reserve(name, 'params')
    if self.has_variable('params', name):
      value = self.get_variable('params', name)
      # Validate that the shape of the init_fn output is the same as the shape
      # of the existing parameter. This is to make sure that the hparams set up
      # in a Flax Module match the shapes coming in during apply, and if not,
      # catch it with an error message.
      # NOTE: We could consider moving this to `self.`
      abs_value = jax.eval_shape(
        lambda: init_fn(random.key(0), *init_args, **init_kwargs)
      )
      abs_value_flat = jax.tree_util.tree_leaves(abs_value)
      value_flat = jax.tree_util.tree_leaves(value)
      for val, abs_val in zip(value_flat, abs_value_flat):
        # NOTE: We could check dtype consistency here as well but it's
        # usefuleness is less obvious. We might intentionally change the dtype
        # for inference to a half float type for example.
        if jnp.shape(val) != jnp.shape(abs_val):
          raise errors.ScopeParamShapeError(
            name, self.path_text, jnp.shape(abs_val), jnp.shape(val)
          )
    else:
      if not self.is_mutable_collection('params'):
        if self.is_collection_empty('params'):
          raise errors.ScopeCollectionNotFound('params', name, self.path_text)
        raise errors.ScopeParamNotFoundError(name, self.path_text)
      value = init_fn(self.make_rng('params'), *init_args, **init_kwargs)
      self.put_variable('params', name, value)
    if unbox:
      value = meta.unbox(value)
    return value

  def _populate_collections(self):
    collections = self.root._variables.keys()  # pylint: disable=protected-access
    for col in collections:
      self._collection(col)

  def has_flag(self, key) -> bool:
    return key in self.flags

  def get_flag(self, key, default=no_flag) -> Any:
    if key not in self.flags and default is no_flag:
      return ValueError(f'Flag {key} not present on scope.')
    return self.flags.get(key, default)


def _unfreeze_variables(variables, mutable):
  new_variables = {}
  for key, value in variables.items():
    if in_filter(mutable, key):
      new_variables[key] = unfreeze(value)
    else:
      new_variables[key] = value
  return new_variables


def bind(
  variables: VariableDict,
  rngs: Optional[RNGSequences] = None,
  mutable: CollectionFilter = False,
  flags: Optional[Mapping] = None,
):
  """Binds variables and rngs to a new ``Scope``.

  bind provides a ``Scope`` instance without transforming a function with
  ``apply``. This is particularly useful for debugging and interactive use cases
  like notebooks where a function would limit the ability split up code into
  different cells.

  a ``Scope`` instance is a stateful object. Note that idiomatic JAX is
  functional and therefore a ``Scope` does not mix well well with vanilla JAX
  APIs. Therefore, we recommend using ``apply`` when code should be reusable and
  compatible across the JAX software ecosystem.

  Args:
    variables: Variable dictionary to bind.
    rngs: RNGs to bind.
    mutable: Which variable collections to treat as mutable.
    flags: internal flags.

  Returns:
    A new scope with the variables and rngs bound to it.
  """
  if not _is_valid_variables(variables):
    raise errors.ApplyScopeInvalidVariablesTypeError()
  if rngs is not None and not _is_valid_rngs(rngs):
    raise errors.InvalidRngError(
      'rngs should be a dictionary mapping strings to `jax.PRNGKey`.'
    )
  new_variables = _unfreeze_variables(variables, mutable)
  return Scope(new_variables, rngs=rngs, mutable=mutable, flags=flags)


def apply(
  fn: Callable[..., Any],
  mutable: CollectionFilter = False,
  flags: Optional[Mapping] = None,
) -> Callable[..., Any]:
  """Functionalize a `Scope` function.

  Args:
    fn: a function taking a `Scope` as its first argument.
    mutable: the filter determining which variable collections are mutable.
    flags: internal flags.

  Returns:
    `fn` with the scope partially applied.
  """

  @functools.wraps(fn)
  def wrapper(
    variables: VariableDict,
    *args,
    rngs: Optional[Union[PRNGKey, RNGSequences]] = None,
    **kwargs,
  ) -> Union[Any, Tuple[Any, Union[VariableDict, Dict[str, Any]]]]:
    if rngs is not None:
      if not _is_valid_rng(rngs) and not _is_valid_rngs(rngs):
        raise ValueError(
          'The ``rngs`` argument passed to an apply function should be a '
          '``jax.PRNGKey`` or a dictionary mapping strings to '
          '``jax.PRNGKey``.'
        )
      if not isinstance(rngs, (dict, FrozenDict)):
        rngs = {'params': rngs}

    # Try to detect if user accidentally passed {'params': {'params': ...}.
    if (
      'params' in variables
      and isinstance(variables['params'], (dict, FrozenDict))
      and 'params' in variables['params']
    ):
      raise errors.ApplyScopeInvalidVariablesStructureError(variables)

    with bind(
      variables, rngs=rngs, mutable=mutable, flags=flags
    ).temporary() as root:
      y = fn(root, *args, **kwargs)
    if mutable is not False:
      return y, root.mutable_variables()
    else:
      return y

  return wrapper


def init(
  fn: Callable[..., Any],
  mutable: CollectionFilter = True,
  flags: Optional[Mapping] = None,
) -> Callable[..., Any]:
  """Functionalize a `Scope` function for initialization.

  Args:
    fn: a function taking a `Scope` as its first argument.
    mutable: the filter determining which variable collections are mutable.
    flags: internal flags.

  Returns:
    `fn` with the scope partially applied.
  """

  @functools.wraps(fn)
  def wrapper(rngs, *args, **kwargs) -> Tuple[Any, VariableDict]:
    if not _is_valid_rng(rngs) and not _is_valid_rngs(rngs):
      raise ValueError(
        'First argument passed to an init function should be a '
        '``jax.PRNGKey`` or a dictionary mapping strings to '
        '``jax.PRNGKey``.'
      )
    if not isinstance(rngs, (dict, FrozenDict)):
      rngs = {'params': rngs}
    init_flags = {**(flags if flags is not None else {}), 'initializing': True}
    return apply(fn, mutable=mutable, flags=init_flags)(
      {}, *args, rngs=rngs, **kwargs
    )

  return wrapper


def lazy_init(
  fn: Callable[..., Any],
  mutable: CollectionFilter = True,
  flags: Optional[Mapping] = None,
) -> Callable[..., Any]:
  """Functionalizes a `Scope` function for lazy initialization.

  Similair to ``init`` except that the init function now accepts
  ``jax.ShapeDtypeStruct`` instances for arguments that do not
  affect the variable initialization (typically this is all the input data).

  Example::

    def f(scope, x):
        # the kernel init only uses the shape of x so we don't actually
        # need a value for x and can pass it as a ShapeDtypeStruct in lazy_init.
        k = scope.param("kernel", nn.initializers.lecun_normal(), (x.shape[-1], x.shape[-1]))
        return x @ k
    init_fn = lazy_init(f)
    variables = init_fn(random.key(0), jax.ShapeDtypeStruct((1, 128), jnp.float32))


  Args:
    fn: a function taking a `Scope` as its first argument.
    mutable: the filter determining which variable collections are mutable.
    flags: internal flags.

  Returns:
    `fn` with the scope partially applied. Unlike ``init`` which returns a tuple of function
    output and variables, the lazy init function only returns the variables.
  """
  return partial_eval.lazy_init(
    lambda *args, **kwargs: init(fn, mutable, flags)(*args, **kwargs)[1]
  )


def _is_valid_collection(col: VariableDict):
  if not isinstance(col, (FrozenDict, dict)):
    return False
  for name in col.keys():
    # Any value can be stored in a collection so only keys can be verified.
    if not isinstance(name, str):
      return False
  return True


def _is_valid_variables(variables: VariableDict) -> bool:
  """Checks whether the given variable dict is valid.

  Args:
    variables: A variable dict.

  Returns:
    True if `variables` is a valid variable dict.
  """
  for name, col in variables.items():
    if not isinstance(name, str):
      return False
    if not _is_valid_collection(col):
      return False
  return True


def _is_valid_rng(rng: Array):
  """Checks whether rng is a valid JAX PRNGKey, also handling custom prngs."""
  # This check is valid for either new-style or old-style PRNG keys
  if not isinstance(rng, (np.ndarray, jnp.ndarray)):
    return False

  # Handle new-style typed PRNG keys
  if hasattr(jax.dtypes, 'prng_key'):  # JAX 0.4.14 or newer
    if jax.dtypes.issubdtype(rng.dtype, jax.dtypes.prng_key):
      return rng.shape == ()
  elif hasattr(jax.random, 'PRNGKeyArray'):  # Previous JAX versions
    if isinstance(rng, jax.random.PRNGKeyArray):
      return rng.shape == ()

  # Handle old-style raw PRNG keys
  expected_rng = jax.eval_shape(
    lambda s: jax.random.key_data(jax.random.key(s)), 0
  )
  if (rng.shape, rng.dtype) != (expected_rng.shape, expected_rng.dtype):
    return False
  return True


def _is_valid_rngs(rngs: Union[PRNGKey, RNGSequences]):
  if not isinstance(rngs, (FrozenDict, dict)):
    return False
  for key, val in rngs.items():
    if not isinstance(key, str):
      return False
    if not _is_valid_rng(val):
      return False
  return True
