code: def __call__(
      self, *args: Any, rngs: tp.Optional[Rngs] = None, **kwargs: Any
  ) -> Any:
      _rngs = (
          {name: stream.key for name, stream in rngs._rngs.items()}
          if rngs
          else {}
      )
  
      variables = {
          collection: value for collection, value in self.states.items()
      }
      out = self.module.apply(variables, *args, rngs=_rngs, **kwargs)
  
      if kwargs.get("mutable", False) != False:
          out, updates = out
          for collection, value in updates.items():
              if collection in self.states:
                  self.states[collection] = value
              else:
                  self.states[collection] = variableslib.variable_type(
                      collection
                  )(value)
  
      return out
  
file_path: flax/flax/experimental/nnx/nnx/compatibility.py
imports: import dataclasses
  import typing
  from typing import Any
  from flax import linen
  from flax.experimental.nnx.nnx import helpers
  from flax.experimental.nnx.nnx import variables
  from flax.experimental.nnx.nnx.module import GraphDef, Module
  from flax.experimental.nnx.nnx.rnglib import Rngs
  from flax.experimental.nnx.nnx.state import State
label_desc: Use image node_id 2 for calling the LinenWrapper obj's underlying member method code with example usage: obj.__call__() and returns: out