class_docstr: A Flax ``nn.Module`` that runs a Haiku transformed function.
  
  This type is designed to make it easy to take a Haiku transformed function
  and/or a Haiku module and use it inside a program that otherwise uses Flax.
  
  Given a Haiku transformed function
  
  >>> def f(x):
  ...   return hk.Linear(1)(x)
  >>> f = hk.transform(f)
  
  You can convert it into a Flax module using:
  
  >>> mod = hk.experimental.flax.Module(f)
  
  Calling this module is the same as calling any regular Flax module:
  
  >>> rng = jax.random.PRNGKey(42)
  >>> x = jnp.ones([1, 1])
  >>> variables = mod.init(rng, x)
  >>> out = mod.apply(variables, x)
  
  If you just want to convert a Haiku module class such that it can be used
  with Flax you can use the ``create`` class method:
  
  >>> mod = hk.experimental.flax.Module.create(hk.Linear, 1)
  >>> variables = mod.init(rng, x)
  >>> out = mod.apply(variables, x)
code: def __call__(self, *args, **kwargs):
      if self.is_initializing():
          rng = self.make_rng("params")
          params, state = self.transformed.init(rng, *args, **kwargs)
          store_haiku_collections(
              self.scope, params=params, state=state
          )
      else:
          params = utils.flatten_flax_to_haiku(
              self.variables.get("params", {})
          )
          state = utils.flatten_flax_to_haiku(
              self.variables.get("state", {})
          )
  
      rng = self.make_rng("apply") if self.has_rng("apply") else None
      out, state_out = self.transformed.apply(
          params, state, rng, *args, **kwargs
      )
      if not self.is_initializing():
          updated_state = only_changed_state(state, state_out)
          store_haiku_collections(self.scope, state=updated_state)
  
      return out
  
file_path: dm-haiku/haiku/_src/flax/flax_module.py
imports: from typing import TypeVar, Union
  import flax.core
  import flax.linen
  from haiku._src import filtering
  from haiku._src import transform
  from haiku._src import typing
  from haiku._src.flax import utils
label_desc: Use image node_id 3 for calling the Module obj's underlying member method code with example usage: obj.__call__() and returns: out