# 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):


import tensorflow as tf
from dreamer.tools import nested
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




def chunk_sequence(sequence, chunk_length, randomize=True, num_chunks=None):
  if 'length' in sequence:
    length = sequence.pop('length')
  else:
    length = tf.shape(nested.flatten(sequence)[0])[0]
  if randomize:
    if not num_chunks:
      num_chunks = tf.maximum(1, length // chunk_length - 1)
    else:
      num_chunks = num_chunks + 0 * length
    used_length = num_chunks * chunk_length
    max_offset = length - used_length
    offset = tf.random_uniform((), 0, max_offset + 1, dtype=tf.int32)
  else:
    if num_chunks is None:
      num_chunks = length // chunk_length
    else:
      num_chunks = num_chunks + 0 * length
    used_length = num_chunks * chunk_length
    offset = 0

# 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/chunk_sequence.py
import tensorflow as tf
from dreamer.tools import nested

# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




def chunk_sequence(sequence, chunk_length, randomize=True, num_chunks=None):
  if 'length' in sequence:
    length = sequence.pop('length')
  else:
    length = tf.shape(nested.flatten(sequence)[0])[0]
  if randomize:
    if not num_chunks:
      num_chunks = tf.maximum(1, length // chunk_length - 1)
    else:
      num_chunks = num_chunks + 0 * length
    used_length = num_chunks * chunk_length
    max_offset = length - used_length
    offset = tf.random_uniform((), 0, max_offset + 1, dtype=tf.int32)
  else:
    if num_chunks is None:
      num_chunks = length // chunk_length
    else:
      num_chunks = num_chunks + 0 * length
    used_length = num_chunks * chunk_length
    offset = 0

  clipped = nested.map(