class_docstr: Base class for Composition bijectors (Chain, JointMap).
  
  A Composition represents a partially ordered set of invertible
  transformations. These transformations may happen in series (Chain), in
  parallel (JointMap), or they could be an arbitrary DAG. Composition handles
  the common machinery of such transformations, delegating graph-traversal to
  `_walk_forward` and `_walk_inverse` (which must be overridden by subclasses).
  
  The `_walk_{direction}` methods take a `step_fn`, a single (structured)
  `argument` (representing zipped `*args`), and arbitrary `**kwargs`. They are
  responsible for invoking `step_fn(bij, bij_inputs, **bij_kwds)`
  for each nested bijector. See `Chain` and `JointMap` for examples.
  
  These methods are typically invoked using `_call_walk_{direction}`, which
  wraps `step_fn` and converts structured `*args` into a single structure of
  tuples, allowing users to provide a `step_fn` with multiple positional
  arguments (e.g., `foward_log_det_jacobian`).
  
  In practice, Bijector methods are defined in the base-class, and users
  should not need to invoke `walk` methods directly.
docstr: Prepares args and calls `_walk_forward`.
  
  Converts a tuple of structured positional arguments to a structure of
  argument tuples, and wraps `step_fn` to unpack inputs and re-pack
  returned values. This way, users may invoke walks using `map_structure`
  semantics, and the concrete `_walk` implementations can operate on
  a single structure of inputs (without worrying about tuple unpacking).
  
  
  For example, the `forward` method looks roughly like:
  ```python
  
  MyComposition()._call_walk_forward(
      lambda bij, x, **kwargs: bij.forward(x, **kwargs),
      composite_inputs, **composite_kwargs)
  ```
  
  More complex methods may need to mutate external state from `step_fn`:
  ```python
  
  shape_trace = {}
  
  def trace_step(bijector, x_shape):
    shape_trace[bijector.name] = x_shape
    return bijector.forward_event_shape(x_shape)
  
  # Calling this populates the `shape_trace` dictionary
  composition.walk_forward(trace_step, composite_input_shape)
  ```
  
  Args:
    step_fn: Callable applied to each wrapped bijector.
      Must accept a bijector instance followed by `len(args)` positional
      arguments whose structures match `bijector.forward_min_event_ndims`,
      and return `len(args)` structures matching
      `bijector.inverse_min_event_ndims`.
    *args: Input arguments propagated to nested bijectors.
    **kwargs: Keyword arguments forwarded to `_walk_forward`.
  Returns:
    The transformed output. If multiple positional arguments are provided, a
      tuple of matching length will be returned.
docstr_tok: ["Prepares","args","and","calls","`","_walk_forward","`",".","Converts","a","tuple","of","structured","positional","arguments","to","a","structure","of","argument","tuples",",","and","wraps","`","step_fn","`","to","unpack","inputs","and","re-pack","returned","values",".","This","way",",","users","may","invoke","walks","using","`","map_structure","`","semantics",",","and","the","concrete","`","_walk","`","implementations","can","operate","on","a","single","structure","of","inputs","(","without","worrying","about","tuple","unpacking",")",".","For","example",",","the","`","forward","`","method","looks","roughly","like",":","``","`","python","MyComposition","(",")","._call_walk_forward","(","lambda","bij",",","x",",","*","*","kwargs",":","bij.forward","(","x",",","*","*","kwargs",")",",","composite_inputs",",","*","*","composite_kwargs",")","``","`","More","complex","methods","may","need","to","mutate","external","state","from","`","step_fn","`",":","``","`","python","shape_trace","=","{","}","def","trace_step","(","bijector",",","x_shape",")",":","shape_trace","[","bijector.name","]","=","x_shape","return","bijector.forward_event_shape","(","x_shape",")","#","Calling","this","populates","the","`","shape_trace","`","dictionary","composition.walk_forward","(","trace_step",",","composite_input_shape",")","``","`","Args",":","step_fn",":","Callable","applied","to","each","wrapped","bijector",".","Must","accept","a","bijector","instance","followed","by","`","len","(","args",")","`","positional","arguments","whose","structures","match","`","bijector.forward_min_event_ndims","`",",","and","return","`","len","(","args",")","`","structures","matching","`","bijector.inverse_min_event_ndims","`",".","*","args",":","Input","arguments","propagated","to","nested","bijectors",".","*","*","kwargs",":","Keyword","arguments","forwarded","to","`","_walk_forward","`",".","Returns",":","The","transformed","output",".","If","multiple","positional","arguments","are","provided",",","a","tuple","of","matching","length","will","be","returned","."]
code: def _call_walk_forward(self, step_fn, *args, **kwargs):
      """Prepares args and calls `_walk_forward`.
  
      Converts a tuple of structured positional arguments to a structure of
      argument tuples, and wraps `step_fn` to unpack inputs and re-pack
      returned values. This way, users may invoke walks using `map_structure`
      semantics, and the concrete `_walk` implementations can operate on
      a single structure of inputs (without worrying about tuple unpacking).
  
  
      For example, the `forward` method looks roughly like:
      ```python
  
      MyComposition()._call_walk_forward(
          lambda bij, x, **kwargs: bij.forward(x, **kwargs),
          composite_inputs, **composite_kwargs)
      ```
  
      More complex methods may need to mutate external state from `step_fn`:
      ```python
  
      shape_trace = {}
  
      def trace_step(bijector, x_shape):
        shape_trace[bijector.name] = x_shape
        return bijector.forward_event_shape(x_shape)
  
      # Calling this populates the `shape_trace` dictionary
      composition.walk_forward(trace_step, composite_input_shape)
      ```
  
      Args:
        step_fn: Callable applied to each wrapped bijector.
          Must accept a bijector instance followed by `len(args)` positional
          arguments whose structures match `bijector.forward_min_event_ndims`,
          and return `len(args)` structures matching
          `bijector.inverse_min_event_ndims`.
        *args: Input arguments propagated to nested bijectors.
        **kwargs: Keyword arguments forwarded to `_walk_forward`.
      Returns:
        The transformed output. If multiple positional arguments are provided, a
          tuple of matching length will be returned.
      """
      args = tuple(
          nest_util.coerce_structure(self.forward_min_event_ndims, x)
          for x in args
      )
  
      if len(args) == 1:
          return self._walk_forward(step_fn, *args, **kwargs)
  
      # Convert a tuple of structures to a structure of tuples. This
      # allows `_walk` methods to route aligned structures of inputs/outputs
      # independently, obviates the need for conditional tuple unpacking.
      packed_args = pack_structs_like(
          self.forward_min_event_ndims, *args
      )
  
      def transform_wrapper(bij, packed_xs, **nested):
          xs = unpack_structs_like(
              bij.forward_min_event_ndims, packed_xs
          )
          ys = step_fn(bij, *xs, **nested)
          return pack_structs_like(bij.inverse_min_event_ndims, *ys)
  
      packed_result = self._walk_forward(
          transform_wrapper, packed_args, **kwargs
      )
      return unpack_structs_like(
          self.inverse_min_event_ndims, packed_result
      )
  
imports: import collections
  import functools
  import sys
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.bijectors import ldj_ratio
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 8 for calling the Composition obj's underlying member method code with example usage: obj._call_walk_forward(step_fn) and returns: unpack_structs_like, self, pack_structs_like