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)
docstr: Converts a given Haiku module into a Flax ``nn.Module``.
  
  TODO(tomhennigan): Support multiple forward methods.
  
  Example usage:
  
  >>> mod = hk.experimental.flax.Module.create(hk.Linear, 1)  # hk.Linear(1)
  >>> rng = jax.random.PRNGKey(42)
  >>> x = jnp.ones([1, 1])
  >>> variables = mod.init(rng, x)
  >>> out = mod.apply(variables, x)
  
  For a stateful module like resnet, you need to also handle output state:
  
  >>> mod = hk.experimental.flax.Module.create(hk.nets.ResNet50, 10)
  >>> rng = jax.random.PRNGKey(42)
  >>> x = jnp.ones([1, 224, 224, 3])
  >>> variables = mod.init(rng, x, is_training=True)
  >>> out, state_out = mod.apply(variables, x, is_training=True,
  ...                            mutable=['state'])
  
  Args:
    hk_cls: A Haiku module type (e.g. ``hk.Linear``).
    *init_args: Positional arguments for the constructor.
    **init_kwargs: Keyword arguments for the constructor.
  
  Returns:
    A Flax ``nn.Module`` wrapping the given class.
docstr_tok: ["Converts","a","given","Haiku","module","into","a","Flax","``","nn.Module","``",".","TODO","(","tomhennigan",")",":","Support","multiple","forward","methods",".","Example","usage",":",">",">",">","mod","=","hk.experimental.flax.Module.create","(","hk.Linear",",","1",")","#","hk.Linear","(","1",")",">",">",">","rng","=","jax.random.PRNGKey","(","42",")",">",">",">","x","=","jnp.ones","(","[","1",",","1","]",")",">",">",">","variables","=","mod.init","(","rng",",","x",")",">",">",">","out","=","mod.apply","(","variables",",","x",")","For","a","stateful","module","like","resnet",",","you","need","to","also","handle","output","state",":",">",">",">","mod","=","hk.experimental.flax.Module.create","(","hk.nets.ResNet50",",","10",")",">",">",">","rng","=","jax.random.PRNGKey","(","42",")",">",">",">","x","=","jnp.ones","(","[","1",",","224",",","224",",","3","]",")",">",">",">","variables","=","mod.init","(","rng",",","x",",","is_training=True",")",">",">",">","out",",","state_out","=","mod.apply","(","variables",",","x",",","is_training=True",",","...","mutable=","[","'state","'","]",")","Args",":","hk_cls",":","A","Haiku","module","type","(","e.g",".","``","hk.Linear","``",")",".","*","init_args",":","Positional","arguments","for","the","constructor",".","*","*","init_kwargs",":","Keyword","arguments","for","the","constructor",".","Returns",":","A","Flax","``","nn.Module","``","wrapping","the","given","class","."]
code: def create(
      cls, hk_cls: type[hk.SupportsCall], *init_args, **init_kwargs
  ) -> "Module":
      """Converts a given Haiku module into a Flax ``nn.Module``.
  
      TODO(tomhennigan): Support multiple forward methods.
  
      Example usage:
  
      >>> mod = hk.experimental.flax.Module.create(hk.Linear, 1)  # hk.Linear(1)
      >>> rng = jax.random.PRNGKey(42)
      >>> x = jnp.ones([1, 1])
      >>> variables = mod.init(rng, x)
      >>> out = mod.apply(variables, x)
  
      For a stateful module like resnet, you need to also handle output state:
  
      >>> mod = hk.experimental.flax.Module.create(hk.nets.ResNet50, 10)
      >>> rng = jax.random.PRNGKey(42)
      >>> x = jnp.ones([1, 224, 224, 3])
      >>> variables = mod.init(rng, x, is_training=True)
      >>> out, state_out = mod.apply(variables, x, is_training=True,
      ...                            mutable=['state'])
  
      Args:
        hk_cls: A Haiku module type (e.g. ``hk.Linear``).
        *init_args: Positional arguments for the constructor.
        **init_kwargs: Keyword arguments for the constructor.
  
      Returns:
        A Flax ``nn.Module`` wrapping the given class.
      """
  
      def fn(*args, **kwargs):
          mod = hk_cls(*init_args, **init_kwargs)
          return mod(*args, **kwargs)
  
      fn = hk.transform_with_state(fn)
      return Module(fn)
  
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 2 for calling the Module obj's underlying member method code with example usage: obj.create(cls, hk_cls) and returns: Module, mod