# 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
      posterior)
  prev_state = nested.map(
      lambda tensor: _merge_dims(tensor, [0, 1]), prev_state)
  (priors, _), _ = tf.nn.dynamic_rnn(
      cell, (obs, sequences['prev_action'], use_obs),
      merged_length,
      prev_state)

  # Restore batch dimension.
  target, prior, posterior, mask = nested.map(
      functools.partial(_restore_batch_dim, batch_size=shape.shape(length)[0]),
      (sequences['target'], priors, sequences['posterior'], sequences['mask']))

  mask = tf.cast(mask, tf.bool)
  return target, prior, posterior, mask


def _merge_dims(tensor, dims):
  if isinstance(tensor, (list, tuple, dict)):
    return nested.map(tensor, lambda x: _merge_dims(x, dims))
  tensor = tf.convert_to_tensor(tensor)
  if (np.array(dims) - min(dims) != np.arange(len(dims))).all():
    raise ValueError('Dimensions to merge must all follow each other.')
  start, end = dims[0], dims[-1]
  output = tf.reshape(tensor, tf.concat([
      tf.shape(tensor)[:start],
      [tf.reduce_prod(tf.shape(tensor)[start: end + 1])],
      tf.shape(tensor)[end + 1:]], axis=0))
  merged = tensor.shape[start: end + 1].as_list()
  output.set_shape(

# 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

      posterior)
  prev_state = nested.map(
      lambda tensor: _merge_dims(tensor, [0, 1]), prev_state)
  (priors, _), _ = tf.nn.dynamic_rnn(
      cell, (obs, sequences['prev_action'], use_obs),
      merged_length,
      prev_state)

  # Restore batch dimension.
  target, prior, posterior, mask = nested.map(
      functools.partial(_restore_batch_dim, batch_size=shape.shape(length)[0]),
      (sequences['target'], priors, sequences['posterior'], sequences['mask']))

  mask = tf.cast(mask, tf.bool)
  return target, prior, posterior, mask


def _merge_dims(tensor, dims):
  if isinstance(tensor, (list, tuple, dict)):
    return nested.map(tensor, lambda x: _merge_dims(x, dims))
  tensor = tf.convert_to_tensor(tensor)
  if (np.array(dims) - min(dims) != np.arange(len(dims))).all():
    raise ValueError('Dimensions to merge must all follow each other.')
  start, end = dims[0], dims[-1]
  output = tf.reshape(tensor, tf.concat([
      tf.shape(tensor)[:start],
      [tf.reduce_prod(tf.shape(tensor)[start: end + 1])],
      tf.shape(tensor)[end + 1:]], axis=0))
  merged = tensor.shape[start: end + 1].as_list()
  output.set_shape(

      tensor.shape[:start].as_list() +