# Path: dreamer/tools/nested.py
# def zip_(*structures, **kwargs):
# def map_(function, *structures, **kwargs):
#   def impl(function, *structures):
# def flatten_(structure):
# def filter_(predicate, *structures, **kwargs):
#   def impl(predicate, *structures):
# 
# Path: dreamer/tools/shape.py
# def shape(tensor):
#   static = tensor.get_shape().as_list()
#   dynamic = tf.unstack(tf.shape(tensor))
#   assert len(static) == len(dynamic)
#   combined = [d if s is None else s for s, d in zip(static, dynamic)]
#   return combined


import functools
import numpy as np
import tensorflow as tf
from dreamer.tools import nested
from dreamer.tools import shape

  progress_fn = lambda tensor: tf.concat([tensor[:, 1:], 0 * tensor[:, :1]], 1)
  other_outputs = tf.scan(
      lambda past_output, _: nested.map(progress_fn, past_output),
      tf.range(amount), first_output)
  sequences = nested.map(
      lambda lhs, rhs: tf.concat([lhs[None], rhs], 0),
      first_output, other_outputs)

  # Merge batch and time dimensions of steps to compute unrolls from every
  # time step as one batch. The time dimension becomes the number of
  # overshooting distances.
  sequences = nested.map(
      lambda tensor: _merge_dims(tensor, [1, 2]),
      sequences)
  sequences = nested.map(
      lambda tensor: tf.transpose(
          tensor, [1, 0] + list(range(2, tensor.shape.ndims))),
      sequences)
  merged_length = tf.reduce_sum(sequences['mask'], 1)

  # Mask out padding frames; unnecessary if the input is already masked.
  sequences = nested.map(
      lambda tensor: tensor * tf.cast(
          _pad_dims(sequences['mask'], tensor.shape.ndims),
          tensor.dtype),
      sequences)

  # Compute open loop rollouts.
  use_obs = tf.zeros(tf.shape(sequences['mask']), tf.bool)[..., None]

# Path: dreamer/tools/nested.py
# def zip_(*structures, **kwargs):
# def map_(function, *structures, **kwargs):
#   def impl(function, *structures):
# def flatten_(structure):
# def filter_(predicate, *structures, **kwargs):
#   def impl(predicate, *structures):
# 
# Path: dreamer/tools/shape.py
# def shape(tensor):
#   static = tensor.get_shape().as_list()
#   dynamic = tf.unstack(tf.shape(tensor))
#   assert len(static) == len(dynamic)
#   combined = [d if s is None else s for s, d in zip(static, dynamic)]
#   return combined


# Path: dreamer/tools/overshooting.py
import functools
import numpy as np
import tensorflow as tf
from dreamer.tools import nested
from dreamer.tools import shape


  progress_fn = lambda tensor: tf.concat([tensor[:, 1:], 0 * tensor[:, :1]], 1)
  other_outputs = tf.scan(
      lambda past_output, _: nested.map(progress_fn, past_output),
      tf.range(amount), first_output)
  sequences = nested.map(
      lambda lhs, rhs: tf.concat([lhs[None], rhs], 0),
      first_output, other_outputs)

  # Merge batch and time dimensions of steps to compute unrolls from every
  # time step as one batch. The time dimension becomes the number of
  # overshooting distances.
  sequences = nested.map(
      lambda tensor: _merge_dims(tensor, [1, 2]),
      sequences)
  sequences = nested.map(
      lambda tensor: tf.transpose(
          tensor, [1, 0] + list(range(2, tensor.shape.ndims))),
      sequences)
  merged_length = tf.reduce_sum(sequences['mask'], 1)

  # Mask out padding frames; unnecessary if the input is already masked.
  sequences = nested.map(
      lambda tensor: tensor * tf.cast(
          _pad_dims(sequences['mask'], tensor.shape.ndims),
          tensor.dtype),
      sequences)

  # Compute open loop rollouts.
  use_obs = tf.zeros(tf.shape(sequences['mask']), tf.bool)[..., None]

  embed_size = nested.flatten(embedded)[0].shape[2].value