class_docstr: Module to compose parameterized functions, implemented as a scan.
class_docstr_tok: ["Module","to","compose","parameterized","functions",",","implemented","as","a","scan","."]
code: def __call__(self, x, *args_ys, reverse=False):
      count = self._count
      init_fn, apply_fn = transform.transform(self._call_wrapped)
  
      def per_layer_init_fn(c, a):
          c, rng = c
          if rng is not None:
              rng, next_rng, apply_rng = jax.random.split(rng, 3)
          else:
              rng, next_rng, apply_rng = None, None, None
          params = init_fn(rng, c, *a)
          c, _ = apply_fn(params, apply_rng, c, *a)
          return (c, next_rng), params
  
      def scanned_init_fn(x, rng):
          _, params = jax.lax.scan(
              per_layer_init_fn, (x, rng), args_ys, length=self._count
          )
          if self._transparency_map is not None:
              return _split_params(
                  params, self._count, self._transparency_map
              )
          else:
              return params
  
      rng = base.maybe_next_rng_key()
  
      try:
          if self._transparency_map is not None:
              lifted_init_fn = lift.transparent_lift(
                  scanned_init_fn, allow_reuse=True
              )
          else:
              lifted_init_fn = lift.lift(
                  scanned_init_fn, allow_reuse=True, name=self._name
              )
          params = lifted_init_fn(x, rng)
      except base.NonEmptyStateError as e:
          raise LayerStackStateError(
              "LayerStack can only be used on Haiku "
              "functions which do not make use of Haiku "
              "state."
          ) from e
  
      # Use scan during apply, threading through random seed so that it's
      # unique for each layer.
      def layer(
          carry: LayerStackCarry, scanned: LayerStackScanned
      ) -> tuple[LayerStackCarry, Any]:
          rng = scanned.rng
          params = scanned.params
  
          kwargs = {}
          if self._pass_reverse_to_layer_fn:
              kwargs["reverse"] = reverse
          out_x, z = apply_fn(
              params, rng, carry.x, *scanned.args_ys, **kwargs
          )
          return LayerStackCarry(x=out_x), z
  
      rng = _get_rng_stack(count)
  
      if self._transparency_map is not None:
          params = _stack_params(
              params, self._count, self._transparency_map
          )
  
      carry = LayerStackCarry(x=x)
      scanned = LayerStackScanned(
          params=params, rng=rng, args_ys=args_ys
      )
  
      carry, zs = jax.lax.scan(
          layer,
          carry,
          scanned,
          length=count,
          unroll=self._unroll,
          reverse=reverse,
      )
      return carry.x, zs
  
imports: import collections
  import functools
  import inspect
  from typing import Any, Callable, Optional, Protocol, Union
  from haiku._src import base
  from haiku._src import lift
  from haiku._src import module
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the _LayerStack obj's underlying member method code with example usage: obj.__call__(x) and returns: carry,  zs, c,  next_rng,  params, LayerStackCarry,  z, _split_params, params