jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __call__(self, params):
    example_decoder = tf_example_decoder.TfExampleSegmentationDecoder()
    def _dataset_parser(value):
      """Parse data to a fixed dimension input image and learning targets.

      Args:
        value: A dictionary contains an image and groundtruth annotations.

      Returns:
        A list of the following elements in order:
        image: Image tensor that is preproessed to have normalized value and
          fixed dimension [image_size, image_size, 3]
        label: label tensor of the same spatial dimension as the image.
      """
      with tf.name_scope('parser'):
        data = example_decoder.decode(value)
        image = data['image']
        label = data['labels_class']
        label = tf.cast(label, tf.int32)
        input_processor = SegmentationInputProcessor(image,
                                                     params['image_size'],
                                                     label)
        # The image normalization is identical to Cloud TPU ResNet.
        input_processor.normalize_image()
        if self._is_training and params['input_rand_hflip']:
          input_processor.random_horizontal_flip()
        if self._is_training:
          input_processor.set_training_random_scale_factors(
              params['train_scale_min'], params['train_scale_max'])
        image = input_processor.resize_and_crop_image()

        # Set padding to background (class=0) during training.
        if self._is_training:
          label = input_processor.resize_and_crop_label(0)
        else:
          label = input_processor.resize_and_crop_label(params['ignore_label'])
        if params['use_bfloat16']:
          image = tf.cast(image, dtype=tf.bfloat16)
        return image, label

    batch_size = params['batch_size']

    dataset = tf.data.Dataset.list_files(
        self._file_pattern, shuffle=self._is_training)
    if self._is_training:
      dataset = dataset.repeat()

    def _prefetch_dataset(filename):
      dataset = tf.data.TFRecordDataset(filename).prefetch(1)
      return dataset

    dataset = dataset.apply(
        contrib_data.parallel_interleave(
            _prefetch_dataset, cycle_length=32, sloppy=self._is_training))
    if self._is_training:
      dataset = dataset.shuffle(64)

    dataset = dataset.map(_dataset_parser, num_parallel_calls=64)
    dataset = dataset.prefetch(batch_size)
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset
//  [python]
def make_kl_weight(global_step, anneal_kl_step):
  if anneal_kl_step > 0:
    kl_weight = tf.clip_by_value(
        tf.to_float(global_step) / tf.to_float(anneal_kl_step), 0., 1.)
  else:
    kl_weight = 1.
  tf.summary.scalar("kl_weight", kl_weight)
  return kl_weight
//  [python]
def make_kl_weight(global_step, anneal_kl_step):
  if anneal_kl_step > 0:
    kl_weight = tf.clip_by_value(
        tf.cast(global_step, tf.float32) / tf.cast(anneal_kl_step, tf.float32), 0., 1.)
  else:
    kl_weight = 1.
  tf.summary.scalar("kl_weight", kl_weight)
  return kl_weight
//  [python]
def log_loss(labels, predictions, epsilon=1e-7, scope=None, weights=None):
  """Calculate log losses.

  Same as tf.losses.log_loss except that this returns the individual losses
  instead of passing them into compute_weighted_loss and returning their
  weighted mean. This is useful for eval jobs that report the mean loss. By
  returning individual losses, that mean loss can be the same regardless of
  batch size.

  Args:
    labels: The ground truth output tensor, same dimensions as 'predictions'.
    predictions: The predicted outputs.
    epsilon: A small increment to add to avoid taking a log of zero.
    scope: The scope for the operations performed in computing the loss.
    weights: Weights to apply to labels.

  Returns:
    A `Tensor` representing the loss values.

  Raises:
    ValueError: If the shape of `predictions` doesn't match that of `labels`.
  """
  with tf.name_scope(scope, "log_loss", (predictions, labels)):
    predictions = tf.to_float(predictions)
    labels = tf.to_float(labels)
    predictions.get_shape().assert_is_compatible_with(labels.get_shape())
    losses = -tf.multiply(labels, tf.log(predictions + epsilon)) - tf.multiply(
        (1 - labels), tf.log(1 - predictions + epsilon))
    if weights is not None:
      losses = tf.multiply(losses, weights)

    return losses
//  [python]
def log_loss(labels, predictions, epsilon=1e-7, scope=None, weights=None):
  """Calculate log losses.

  Same as tf.losses.log_loss except that this returns the individual losses
  instead of passing them into compute_weighted_loss and returning their
  weighted mean. This is useful for eval jobs that report the mean loss. By
  returning individual losses, that mean loss can be the same regardless of
  batch size.

  Args:
    labels: The ground truth output tensor, same dimensions as 'predictions'.
    predictions: The predicted outputs.
    epsilon: A small increment to add to avoid taking a log of zero.
    scope: The scope for the operations performed in computing the loss.
    weights: Weights to apply to labels.

  Returns:
    A `Tensor` representing the loss values.

  Raises:
    ValueError: If the shape of `predictions` doesn't match that of `labels`.
  """
  with tf.name_scope(scope, "log_loss", (predictions, labels)):
    predictions = tf.cast(predictions, tf.float32)
    labels = tf.cast(labels, tf.float32)
    predictions.get_shape().assert_is_compatible_with(labels.get_shape())
    losses = -tf.multiply(labels, tf.log(predictions + epsilon)) - tf.multiply(
        (1 - labels), tf.log(1 - predictions + epsilon))
    if weights is not None:
      losses = tf.multiply(losses, weights)

    return losses
//  [python]
def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
                      is_training=False, random_latent=False,
                      tiny_mode=False, small_mode=False):
  """Builds convolutional latent tower for stochastic model.

  At training time this tower generates a latent distribution (mean and std)
  conditioned on the entire video. This latent variable will be fed to the
  main tower as an extra variable to be used for future frames prediction.
  At inference time, the tower is disabled and only returns latents sampled
  from N(0,1).
  If the multi_latent flag is on, a different latent for every timestep would
  be generated.

  Args:
    images: tensor of ground truth image sequences
    time_axis: the time axis  in images tensor
    latent_channels: number of latent channels
    min_logvar: minimum value for log_var
    is_training: whether or not it is training mode
    random_latent: whether or not generate random latents
    tiny_mode: whether or not it is tiny_mode. tiny_mode sets the number
        of conv channels to 1 at each layer. useful for testing the
        integration tests.
    small_mode: whether or not it is small_mode. small mode is the same model
        with less conv and lstm layers and also lower number of channels.
        suitable for videos with less complexity and testing.
  Returns:
    latent_mean: predicted latent mean
    latent_logvar: predicted latent log variance
  """
  conv_size = tinyify([32, 64, 64], tiny_mode, small_mode)
  with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
    images = tf.to_float(images)
    images = tf.unstack(images, axis=time_axis)
    images = tf.concat(images, axis=3)

    x = images
    x = common_layers.make_even_size(x)
    x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                   padding="SAME", activation=tf.nn.relu, name="latent_conv1")
    x = contrib.layers().layer_norm(x)
    if not small_mode:
      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                     padding="SAME", activation=tf.nn.relu, name="latent_conv2")
      x = contrib.layers().layer_norm(x)
    x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                   padding="SAME", activation=tf.nn.relu, name="latent_conv3")
    x = contrib.layers().layer_norm(x)

    nc = latent_channels
    mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
                      padding="SAME", activation=None, name="latent_mean")
    logv = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_std")
    logvar = logv + min_logvar

    # No latent tower at inference time, just standard gaussian.
    if not is_training:
      return tf.zeros_like(mean), tf.zeros_like(logvar)

    # No latent in the first phase
    ret_mean, ret_logvar = tf.cond(
        random_latent,
        lambda: (tf.zeros_like(mean), tf.zeros_like(logvar)),
        lambda: (mean, logvar))

    return ret_mean, ret_logvar
//  [python]
def conv_latent_tower(images, time_axis, latent_channels=1, min_logvar=-5,
                      is_training=False, random_latent=False,
                      tiny_mode=False, small_mode=False):
  """Builds convolutional latent tower for stochastic model.

  At training time this tower generates a latent distribution (mean and std)
  conditioned on the entire video. This latent variable will be fed to the
  main tower as an extra variable to be used for future frames prediction.
  At inference time, the tower is disabled and only returns latents sampled
  from N(0,1).
  If the multi_latent flag is on, a different latent for every timestep would
  be generated.

  Args:
    images: tensor of ground truth image sequences
    time_axis: the time axis  in images tensor
    latent_channels: number of latent channels
    min_logvar: minimum value for log_var
    is_training: whether or not it is training mode
    random_latent: whether or not generate random latents
    tiny_mode: whether or not it is tiny_mode. tiny_mode sets the number
        of conv channels to 1 at each layer. useful for testing the
        integration tests.
    small_mode: whether or not it is small_mode. small mode is the same model
        with less conv and lstm layers and also lower number of channels.
        suitable for videos with less complexity and testing.
  Returns:
    latent_mean: predicted latent mean
    latent_logvar: predicted latent log variance
  """
  conv_size = tinyify([32, 64, 64], tiny_mode, small_mode)
  with tf.variable_scope("latent", reuse=tf.AUTO_REUSE):
    images = tf.cast(images, tf.float32)
    images = tf.unstack(images, axis=time_axis)
    images = tf.concat(images, axis=3)

    x = images
    x = common_layers.make_even_size(x)
    x = tfl.conv2d(x, conv_size[0], [3, 3], strides=(2, 2),
                   padding="SAME", activation=tf.nn.relu, name="latent_conv1")
    x = contrib.layers().layer_norm(x)
    if not small_mode:
      x = tfl.conv2d(x, conv_size[1], [3, 3], strides=(2, 2),
                     padding="SAME", activation=tf.nn.relu, name="latent_conv2")
      x = contrib.layers().layer_norm(x)
    x = tfl.conv2d(x, conv_size[2], [3, 3], strides=(1, 1),
                   padding="SAME", activation=tf.nn.relu, name="latent_conv3")
    x = contrib.layers().layer_norm(x)

    nc = latent_channels
    mean = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
                      padding="SAME", activation=None, name="latent_mean")
    logv = tfl.conv2d(x, nc, [3, 3], strides=(2, 2),
                      padding="SAME", activation=tf.nn.relu, name="latent_std")
    logvar = logv + min_logvar

    # No latent tower at inference time, just standard gaussian.
    if not is_training:
      return tf.zeros_like(mean), tf.zeros_like(logvar)

    # No latent in the first phase
    ret_mean, ret_logvar = tf.cond(
        random_latent,
        lambda: (tf.zeros_like(mean), tf.zeros_like(logvar)),
        lambda: (mean, logvar))

    return ret_mean, ret_logvar
//  [python]
def correct(boxes, keypoints, image_height, image_width):
    """
    Remove boxes and keypoints that are outside of the image.

    Arguments:
        boxes: a float tensor with shape [num_persons, 4].
        keypoints: an int tensor with shape [num_persons, 17, 3].
        image_height, image_width: float tensors with shape [].
    Returns:
        boxes: a float tensor with shape [num_remaining_boxes, 4],
            where num_remaining_boxes <= num_persons.
        keypoints: an int tensor with shape [num_remaining_boxes, 17, 3].
    """

    window = tf.stack([0.0, 0.0, image_height, image_width])
    boxes, keep_indices = prune_non_overlapping_boxes(
        boxes, tf.expand_dims(window, 0),
        min_overlap=OVERLAP_THRESHOLD
    )

    ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=1)
    ymin = tf.clip_by_value(ymin, 0.0, image_height)
    xmin = tf.clip_by_value(xmin, 0.0, image_width)
    ymax = tf.clip_by_value(ymax, 0.0, image_height)
    xmax = tf.clip_by_value(xmax, 0.0, image_width)
    boxes = tf.stack([ymin, xmin, ymax, xmax], axis=1)

    keypoints = tf.gather(keypoints, keep_indices)
    y, x, v = tf.split(keypoints, 3, axis=2)

    image_height = tf.to_int32(image_height)
    image_width = tf.to_int32(image_width)

    coordinate_violations = tf.concat([
        tf.less(y, 0), tf.less(x, 0),
        tf.greater_equal(y, image_height),
        tf.greater_equal(x, image_width)
    ], axis=2)  # shape [num_remaining_boxes, 17, 4]

    valid_indicator = tf.logical_not(tf.reduce_any(coordinate_violations, axis=2))
    valid_indicator = tf.expand_dims(valid_indicator, 2)
    # it has shape [num_remaining_boxes, 17, 1]

    v *= tf.to_int32(valid_indicator)
    keypoints = tf.concat([y, x, v], axis=2)

    return boxes, keypoints
//  [python]
def correct(boxes, keypoints, image_height, image_width):
    """
    Remove boxes and keypoints that are outside of the image.

    Arguments:
        boxes: a float tensor with shape [num_persons, 4].
        keypoints: an int tensor with shape [num_persons, 17, 3].
        image_height, image_width: float tensors with shape [].
    Returns:
        boxes: a float tensor with shape [num_remaining_boxes, 4],
            where num_remaining_boxes <= num_persons.
        keypoints: an int tensor with shape [num_remaining_boxes, 17, 3].
    """

    window = tf.stack([0.0, 0.0, image_height, image_width])
    boxes, keep_indices = prune_non_overlapping_boxes(
        boxes, tf.expand_dims(window, 0),
        min_overlap=OVERLAP_THRESHOLD
    )

    ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=1)
    ymin = tf.clip_by_value(ymin, 0.0, image_height)
    xmin = tf.clip_by_value(xmin, 0.0, image_width)
    ymax = tf.clip_by_value(ymax, 0.0, image_height)
    xmax = tf.clip_by_value(xmax, 0.0, image_width)
    boxes = tf.stack([ymin, xmin, ymax, xmax], axis=1)

    keypoints = tf.gather(keypoints, keep_indices)
    y, x, v = tf.split(keypoints, 3, axis=2)

    image_height = tf.cast(image_height, tf.int32)
    image_width = tf.cast(image_width, tf.int32)

    coordinate_violations = tf.concat([
        tf.less(y, 0), tf.less(x, 0),
        tf.greater_equal(y, image_height),
        tf.greater_equal(x, image_width)
    ], axis=2)  # shape [num_remaining_boxes, 17, 4]

    valid_indicator = tf.logical_not(tf.reduce_any(coordinate_violations, axis=2))
    valid_indicator = tf.expand_dims(valid_indicator, 2)
    # it has shape [num_remaining_boxes, 17, 1]

    v *= tf.cast(valid_indicator, tf.int32)
    keypoints = tf.concat([y, x, v], axis=2)

    return boxes, keypoints
//  [python]
def jitter_scale(image, image_size, keypoints, center, scale_range):
    with tf.name_scope(None, 'jitter_scale', [image, image_size, keypoints]):
        scale_factor = tf.random_uniform(
            [1],
            minval=scale_range[0],
            maxval=scale_range[1],
            dtype=tf.float32)
        new_size = tf.to_int32(tf.to_float(image_size) * scale_factor)
        new_image = tf.image.resize_images(image, new_size)

        # This is [height, width] -> [y, x] -> [col, row]
        actual_factor = tf.to_float(
            tf.shape(new_image)[:2]) / tf.to_float(image_size)
        x = keypoints[0, :] * actual_factor[1]
        y = keypoints[1, :] * actual_factor[0]

        cx = tf.cast(center[0], actual_factor.dtype) * actual_factor[1]
        cy = tf.cast(center[1], actual_factor.dtype) * actual_factor[0]

        return new_image, tf.stack([x, y]), tf.cast(
            tf.stack([cx, cy]), tf.int32)
//  [python]
def jitter_scale(image, image_size, keypoints, center, scale_range):
    with tf.name_scope(None, 'jitter_scale', [image, image_size, keypoints]):
        scale_factor = tf.random_uniform(
            [1],
            minval=scale_range[0],
            maxval=scale_range[1],
            dtype=tf.float32)
        new_size = tf.cast(tf.to_float(image_size) * scale_factor, tf.int32)
        new_image = tf.image.resize_images(image, new_size)

        # This is [height, width] -> [y, x] -> [col, row]
        actual_factor = tf.to_float(
            tf.shape(new_image)[:2]) / tf.to_float(image_size)
        x = keypoints[0, :] * actual_factor[1]
        y = keypoints[1, :] * actual_factor[0]

        cx = tf.cast(center[0], actual_factor.dtype) * actual_factor[1]
        cy = tf.cast(center[1], actual_factor.dtype) * actual_factor[0]

        return new_image, tf.stack([x, y]), tf.cast(
            tf.stack([cx, cy]), tf.int32)
//  [python]
def projective_inverse_warp(
    img, depth, pose, intrinsics, ret_flows=False):
  """Inverse warp a source image to the target image plane based on projection.

  Args:
    img: the source image [batch, height_s, width_s, 3]
    depth: depth map of the target image [batch, height_t, width_t]
    pose: target to source camera transformation matrix [batch, 4, 4]
    intrinsics: camera intrinsics [batch, 3, 3]
    ret_flows: whether to return the displacements/flows as well
  Returns:
    Source image inverse warped to the target image plane [batch, height_t,
    width_t, 3]
  """
  num_depths = tf.shape(depth)[0]
  batch = tf.to_int32(tf.shape(img)[0]/num_depths)
  height = tf.shape(img)[1]
  width = tf.shape(img)[2]
  # Construct pixel grid coordinates
  pixel_coords = meshgrid_abs(batch*num_depths, height, width)
  # Convert pixel coordinates to the camera frame
  cam_coords = pixel2cam(depth, pixel_coords, intrinsics)
  # Construct a 4x4 intrinsic matrix (TODO: can it be 3x4?)
  filler = tf.constant([0.0, 0.0, 0.0, 1.0], shape=[1, 1, 4])
  filler = tf.tile(filler, [batch * num_depths, 1, 1])
  intrinsics = tf.concat(
      [intrinsics, tf.zeros([batch * num_depths, 3, 1])], axis=2)
  intrinsics = tf.concat([intrinsics, filler], axis=1)
  # Get a 4x4 transformation matrix from 'target' camera frame to 'source'
  # pixel frame.
  proj_tgt_cam_to_src_pixel = tf.matmul(intrinsics, pose)
  src_pixel_coords = cam2pixel(cam_coords, proj_tgt_cam_to_src_pixel)
  output_img = contrib_resampler.resampler(img, src_pixel_coords)
  if ret_flows:
    return output_img, src_pixel_coords - cam_coords
  else:
    return output_img
//  [python]
def projective_inverse_warp(
    img, depth, pose, intrinsics, ret_flows=False):
  """Inverse warp a source image to the target image plane based on projection.

  Args:
    img: the source image [batch, height_s, width_s, 3]
    depth: depth map of the target image [batch, height_t, width_t]
    pose: target to source camera transformation matrix [batch, 4, 4]
    intrinsics: camera intrinsics [batch, 3, 3]
    ret_flows: whether to return the displacements/flows as well
  Returns:
    Source image inverse warped to the target image plane [batch, height_t,
    width_t, 3]
  """
  num_depths = tf.shape(depth)[0]
  batch = tf.cast(tf.shape(img)[0]/num_depths, tf.int32)
  height = tf.shape(img)[1]
  width = tf.shape(img)[2]
  # Construct pixel grid coordinates
  pixel_coords = meshgrid_abs(batch*num_depths, height, width)
  # Convert pixel coordinates to the camera frame
  cam_coords = pixel2cam(depth, pixel_coords, intrinsics)
  # Construct a 4x4 intrinsic matrix (TODO: can it be 3x4?)
  filler = tf.constant([0.0, 0.0, 0.0, 1.0], shape=[1, 1, 4])
  filler = tf.tile(filler, [batch * num_depths, 1, 1])
  intrinsics = tf.concat(
      [intrinsics, tf.zeros([batch * num_depths, 3, 1])], axis=2)
  intrinsics = tf.concat([intrinsics, filler], axis=1)
  # Get a 4x4 transformation matrix from 'target' camera frame to 'source'
  # pixel frame.
  proj_tgt_cam_to_src_pixel = tf.matmul(intrinsics, pose)
  src_pixel_coords = cam2pixel(cam_coords, proj_tgt_cam_to_src_pixel)
  output_img = contrib_resampler.resampler(img, src_pixel_coords)
  if ret_flows:
    return output_img, src_pixel_coords - cam_coords
  else:
    return output_img
//  [python]
def tgt_coords(depth_planes_tile, pose, intrinsics):
  """Calculate coordinates for MPI voxels in target camera frame.

  Args:
    depth_planes_tile: depth values for each plane
    pose: source to target camera transformation [B, 4, 4]
    intrinsics: camera intrinsics [B, 3, 3]
  Returns:
    coordinates for MPI voxels in target camera frame
  """
  batch = tf.shape(pose)[0]
  num_depths = tf.to_int32(tf.shape(depth_planes_tile)[0]/batch)
  height = tf.shape(depth_planes_tile)[1]
  width = tf.shape(depth_planes_tile)[2]

  pose_tile = tf.tile(pose, [num_depths, 1, 1])
  intrinsics_tile = tf.tile(intrinsics, [num_depths, 1, 1])

  # Construct pixel grid coordinates
  pixel_coords = meshgrid_abs(batch*num_depths, height, width)
  # Convert pixel coordinates to the camera frame
  cam_coords = pixel2cam(depth_planes_tile, pixel_coords, intrinsics_tile)

  cam_coords = tf.reshape(cam_coords, [batch*num_depths, 4, -1])
  transformed_coords = tf.matmul(pose_tile, cam_coords)
  transformed_coords = tf.reshape(transformed_coords,
                                  [batch, num_depths, 4, height, width])
  transformed_coords = tf.transpose(transformed_coords, [0, 3, 4, 1, 2])
  t_d = transformed_coords[:, :, :, :, 3]
  t_x = transformed_coords[:, :, :, :, 0] / (t_d + 1.0e-10)
  t_y = transformed_coords[:, :, :, :, 1] / (t_d + 1.0e-10)
  t_z = transformed_coords[:, :, :, :, 2] / (t_d + 1.0e-10)
  t_xyz = tf.stack([t_x, t_y, t_z], axis=4)

  return t_xyz
//  [python]
def tgt_coords(depth_planes_tile, pose, intrinsics):
  """Calculate coordinates for MPI voxels in target camera frame.

  Args:
    depth_planes_tile: depth values for each plane
    pose: source to target camera transformation [B, 4, 4]
    intrinsics: camera intrinsics [B, 3, 3]
  Returns:
    coordinates for MPI voxels in target camera frame
  """
  batch = tf.shape(pose)[0]
  num_depths = tf.cast(tf.shape(depth_planes_tile)[0]/batch, tf.int32)
  height = tf.shape(depth_planes_tile)[1]
  width = tf.shape(depth_planes_tile)[2]

  pose_tile = tf.tile(pose, [num_depths, 1, 1])
  intrinsics_tile = tf.tile(intrinsics, [num_depths, 1, 1])

  # Construct pixel grid coordinates
  pixel_coords = meshgrid_abs(batch*num_depths, height, width)
  # Convert pixel coordinates to the camera frame
  cam_coords = pixel2cam(depth_planes_tile, pixel_coords, intrinsics_tile)

  cam_coords = tf.reshape(cam_coords, [batch*num_depths, 4, -1])
  transformed_coords = tf.matmul(pose_tile, cam_coords)
  transformed_coords = tf.reshape(transformed_coords,
                                  [batch, num_depths, 4, height, width])
  transformed_coords = tf.transpose(transformed_coords, [0, 3, 4, 1, 2])
  t_d = transformed_coords[:, :, :, :, 3]
  t_x = transformed_coords[:, :, :, :, 0] / (t_d + 1.0e-10)
  t_y = transformed_coords[:, :, :, :, 1] / (t_d + 1.0e-10)
  t_z = transformed_coords[:, :, :, :, 2] / (t_d + 1.0e-10)
  t_xyz = tf.stack([t_x, t_y, t_z], axis=4)

  return t_xyz
//  [python]
def flow_gather(source_images, flows):
  """Gather from a tensor of images.

  Args:
    source_images: 5D tensor of images [B, H, W, D, 3]
    flows: 5D tensor of x/y offsets to gather for each slice (pixel offsets)
  Returns:
    warped_imgs_reshape: 5D tensor of gathered (warped) images [B, H, W, D, 3]
  """
  batchsize = tf.shape(source_images)[0]
  height = tf.shape(source_images)[1]
  width = tf.shape(source_images)[2]
  num_depths = tf.shape(source_images)[3]
  source_images_reshape = tf.reshape(
      tf.transpose(source_images, [0, 3, 1, 2, 4]),
      [batchsize * num_depths, height, width, 3])
  flows_reshape = tf.reshape(
      tf.transpose(flows, [0, 3, 1, 2, 4]),
      [batchsize * num_depths, height, width, 2])
  _, h, w = tf.meshgrid(
      tf.range(tf.to_float(batchsize * num_depths), dtype=tf.float32),
      tf.range(tf.to_float(height), dtype=tf.float32),
      tf.range(tf.to_float(width), dtype=tf.float32),
      indexing='ij')
  coords_y = tf.clip_by_value(h + flows_reshape[Ellipsis, 0], 0.0,
                              tf.to_float(height))
  coords_x = tf.clip_by_value(w + flows_reshape[Ellipsis, 1], 0.0,
                              tf.to_float(width))
  sampling_coords = tf.stack([coords_x, coords_y], axis=-1)
  warped_imgs = contrib_resampler.resampler(source_images_reshape,
                                            sampling_coords)
  warped_imgs_reshape = tf.transpose(
      tf.reshape(warped_imgs, [batchsize, num_depths, height, width, 3]),
      [0, 2, 3, 1, 4])
  return warped_imgs_reshape
//  [python]
def flow_gather(source_images, flows):
  """Gather from a tensor of images.

  Args:
    source_images: 5D tensor of images [B, H, W, D, 3]
    flows: 5D tensor of x/y offsets to gather for each slice (pixel offsets)
  Returns:
    warped_imgs_reshape: 5D tensor of gathered (warped) images [B, H, W, D, 3]
  """
  batchsize = tf.shape(source_images)[0]
  height = tf.shape(source_images)[1]
  width = tf.shape(source_images)[2]
  num_depths = tf.shape(source_images)[3]
  source_images_reshape = tf.reshape(
      tf.transpose(source_images, [0, 3, 1, 2, 4]),
      [batchsize * num_depths, height, width, 3])
  flows_reshape = tf.reshape(
      tf.transpose(flows, [0, 3, 1, 2, 4]),
      [batchsize * num_depths, height, width, 2])
  _, h, w = tf.meshgrid(
      tf.range(tf.cast(batchsize * num_depths, dtype=tf.float32)),
      tf.range(tf.cast(height, dtype=tf.float32)),
      tf.range(tf.cast(width, dtype=tf.float32)),
      indexing='ij')
  coords_y = tf.clip_by_value(h + flows_reshape[Ellipsis, 0], 0.0,
                              tf.cast(height, dtype=tf.float32))
  coords_x = tf.clip_by_value(w + flows_reshape[Ellipsis, 1], 0.0,
                              tf.cast(width, dtype=tf.float32))
  sampling_coords = tf.stack([coords_x, coords_y], axis=-1)
  warped_imgs = contrib_resampler.resampler(source_images_reshape,
                                            sampling_coords)
  warped_imgs_reshape = tf.transpose(
      tf.reshape(warped_imgs, [batchsize, num_depths, height, width, 3]),
      [0, 2, 3, 1, 4])
  return warped_imgs_reshape
//  [python]
def image_with_colormap(name,
                        tensor,
                        colormap_name,
                        min_value=None,
                        max_value=None):
  """Creates an image summary using a colormap if summaries are enabled.

  Args:
    name: A string, name of the summary.
    tensor: A tensor of rank 3, batch x height x width, from which the images
      are to be created.
    colormap_name: A string, must be one of matplotlib colormaps, see
      https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html
    min_value: A scalar, the value in `tensor`that will be mapped to zero. If
      None, the minimum value of `tensor` per item will be used.
    max_value: A scalar, the value in `tensor`that will be mapped to 1.0. If
      None, the minimum value of `tensor` per item will be used.

  Returns:
    A tensorlfow summary op, or None.
  """
  if tensor.shape.rank != 3:
    raise ValueError('Tensor\'s rank has to be 3, not %s.' %
                     str(tensor.shape.rank))
  if not _summaries_enabled:
    return None
  if min_value is None:
    min_value = tf.reduce_min(tensor, axis=[1, 2], keepdims=True)
  if max_value is None:
    max_value = tf.reduce_max(tensor, axis=[1, 2], keepdims=True)
  normalized_tensor = (tensor - min_value) / (max_value - min_value + 1e-12)
  levels = tf.to_int32(tf.clip_by_value(normalized_tensor, 0.0, 1.0) * 255.0)
  colormap = cm.get_cmap(colormap_name)(range(256))[:, :3]  # Ignore the alpha
  colormapped_image = tf.gather(params=colormap, indices=levels)
  return tf.summary.image(name, colormapped_image)
//  [python]
def image_with_colormap(name,
                        tensor,
                        colormap_name,
                        min_value=None,
                        max_value=None):
  """Creates an image summary using a colormap if summaries are enabled.

  Args:
    name: A string, name of the summary.
    tensor: A tensor of rank 3, batch x height x width, from which the images
      are to be created.
    colormap_name: A string, must be one of matplotlib colormaps, see
      https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html
    min_value: A scalar, the value in `tensor`that will be mapped to zero. If
      None, the minimum value of `tensor` per item will be used.
    max_value: A scalar, the value in `tensor`that will be mapped to 1.0. If
      None, the minimum value of `tensor` per item will be used.

  Returns:
    A tensorlfow summary op, or None.
  """
  if tensor.shape.rank != 3:
    raise ValueError('Tensor\'s rank has to be 3, not %s.' %
                     str(tensor.shape.rank))
  if not _summaries_enabled:
    return None
  if min_value is None:
    min_value = tf.reduce_min(tensor, axis=[1, 2], keepdims=True)
  if max_value is None:
    max_value = tf.reduce_max(tensor, axis=[1, 2], keepdims=True)
  normalized_tensor = (tensor - min_value) / (max_value - min_value + 1e-12)
  levels = tf.cast(tf.clip_by_value(normalized_tensor, 0.0, 1.0) * 255.0, tf.int32)
  colormap = cm.get_cmap(colormap_name)(range(256))[:, :3]  # Ignore the alpha
  colormapped_image = tf.gather(params=colormap, indices=levels)
  return tf.summary.image(name, colormapped_image)
//  [python]
def summarize(self):
        """Summarize the number of positives and negatives after mining."""
        if self._num_positives_list and self._num_negatives_list:
            avg_num_positives = tf.reduce_mean(tf.to_float(self._num_positives_list))
            avg_num_negatives = tf.reduce_mean(tf.to_float(self._num_negatives_list))
            tf.summary.scalar('HardExampleMiner/NumPositives', avg_num_positives)
            tf.summary.scalar('HardExampleMiner/NumNegatives', avg_num_negatives)
//  [python]
def summarize(self):
        """Summarize the number of positives and negatives after mining."""
        if self._num_positives_list and self._num_negatives_list:
            avg_num_positives = tf.reduce_mean(tf.cast(self._num_positives_list, tf.float32))
            avg_num_negatives = tf.reduce_mean(tf.cast(self._num_negatives_list, tf.float32))
            tf.summary.scalar('HardExampleMiner/NumPositives', avg_num_positives)
            tf.summary.scalar('HardExampleMiner/NumNegatives', avg_num_negatives)
//  [python]
def batch_boolean_mask(mask):
  """Get indices of true values.

  Args:
    mask: [batch_size, num_values]

  Returns:
    true_indices: [batch_size, max_true]
    gathered_mask: [batch_size, max_true]
  """
  # [batch_size, num_values]
  mask = tf.to_int32(mask)

  # [batch_size]
  num_true = tf.reduce_sum(mask, 1)

  # []
  max_true = tf.reduce_max(num_true)

  # [batch_size, max_true]
  gathered_mask, true_indices = tf.nn.top_k(mask, max_true)
  gathered_mask = tf.cast(gathered_mask, tf.bool)

  return gathered_mask, true_indices
//  [python]
def batch_boolean_mask(mask):
  """Get indices of true values.

  Args:
    mask: [batch_size, num_values]

  Returns:
    true_indices: [batch_size, max_true]
    gathered_mask: [batch_size, max_true]
  """
  # [batch_size, num_values]
  mask = tf.cast(mask, tf.int32)

  # [batch_size]
  num_true = tf.reduce_sum(mask, 1)

  # []
  max_true = tf.reduce_max(num_true)

  # [batch_size, max_true]
  gathered_mask, true_indices = tf.nn.top_k(mask, max_true)
  gathered_mask = tf.cast(gathered_mask, tf.bool)

  return gathered_mask, true_indices
//  [python]
def get_timing_signal(length,
                      min_timescale=1,
                      max_timescale=1e4,
                      num_timescales=16):
  """Create Tensor of sinusoids of different frequencies.

  Args:
    length: Length of the Tensor to create, i.e. Number of steps.
    min_timescale: a float
    max_timescale: a float
    num_timescales: an int
  Returns:
    Tensor of shape [length, 2 * num_timescales].
  """
  positions = tf.to_float(tf.range(length))
  log_timescale_increment = (
      math.log(max_timescale / min_timescale) / (num_timescales - 1))
  inv_timescales = min_timescale * tf.exp(
      tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
  scaled_time = tf.expand_dims(positions, 1) * tf.expand_dims(inv_timescales, 0)
  return tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
//  [python]
def get_timing_signal(length,
                      min_timescale=1,
                      max_timescale=1e4,
                      num_timescales=16):
  """Create Tensor of sinusoids of different frequencies.

  Args:
    length: Length of the Tensor to create, i.e. Number of steps.
    min_timescale: a float
    max_timescale: a float
    num_timescales: an int
  Returns:
    Tensor of shape [length, 2 * num_timescales].
  """
  positions = tf.cast(tf.range(length), tf.float32)
  log_timescale_increment = (
      math.log(max_timescale / min_timescale) / (num_timescales - 1))
  inv_timescales = min_timescale * tf.exp(
      tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment)
  scaled_time = tf.expand_dims(positions, 1) * tf.expand_dims(inv_timescales, 0)
  return tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
//  [python]
def build_model(self):
        with slim.arg_scope([slim.conv2d, slim.conv2d_transpose], activation_fn=tf.nn.elu):
            with tf.variable_scope('model', reuse=self.reuse_variables):

                self.left_pyramid  = self.scale_pyramid(self.left,  4)
                if self.mode == 'train':
                    self.right_pyramid = self.scale_pyramid(self.right, 4)
                    self.left_invDepth_pyramid = self.scale_pyramid(self.left_invDepth, 4, isDepth=True)
                    self.right_invDepth_pyramid = self.scale_pyramid(self.right_invDepth, 4, isDepth=True)                    
                    self.left_invDepth_mask_pyramid= [tf.to_float(tf.greater(self.left_invDepth_pyramid[i],0.01),name='ToFloat') for i in range(4)]
                    self.right_invDepth_mask_pyramid= [tf.to_float(tf.greater(self.right_invDepth_pyramid[i],0.01),name='ToFloat') for i in range(4)]


                if self.params.do_stereo:
                    self.model_input = tf.concat([self.left, self.right], 3)
                else:
                    self.model_input = self.left

                #build model
                if self.params.encoder == 'vgg':
                    self.build_vgg()
                elif self.params.encoder == 'resnet50':
                    self.build_resnet50()
                elif self.params.encoder == 'vggforward':
                    self.invDepth_left_est= self.build_vgg_forward(self.left,'resnet50-forward',reuse=False)
                    if self.mode == 'train':
                        self.invDepth_right_est= self.build_vgg_forward(self.right,'resnet50-forward',reuse=True)
                else:
                    self.invDepth_left_est= self.build_resnet50_forward(self.left,'resnet50-forward',reuse=False)
                    if self.mode == 'train':
                        self.invDepth_right_est= self.build_resnet50_forward(self.right,'resnet50-forward',reuse=True)
//  [python]
def build_model(self):
        with slim.arg_scope([slim.conv2d, slim.conv2d_transpose], activation_fn=tf.nn.elu):
            with tf.variable_scope('model', reuse=self.reuse_variables):

                self.left_pyramid  = self.scale_pyramid(self.left,  4)
                if self.mode == 'train':
                    self.right_pyramid = self.scale_pyramid(self.right, 4)
                    self.left_invDepth_pyramid = self.scale_pyramid(self.left_invDepth, 4, isDepth=True)
                    self.right_invDepth_pyramid = self.scale_pyramid(self.right_invDepth, 4, isDepth=True)                    
                    self.left_invDepth_mask_pyramid= [tf.cast(tf.greater(self.left_invDepth_pyramid[i],0.01),tf.float32,name='ToFloat') for i in range(4)]
                    self.right_invDepth_mask_pyramid= [tf.cast(tf.greater(self.right_invDepth_pyramid[i],0.01),tf.float32,name='ToFloat') for i in range(4)]


                if self.params.do_stereo:
                    self.model_input = tf.concat([self.left, self.right], 3)
                else:
                    self.model_input = self.left

                #build model
                if self.params.encoder == 'vgg':
                    self.build_vgg()
                elif self.params.encoder == 'resnet50':
                    self.build_resnet50()
                elif self.params.encoder == 'vggforward':
                    self.invDepth_left_est= self.build_vgg_forward(self.left,'resnet50-forward',reuse=False)
                    if self.mode == 'train':
                        self.invDepth_right_est= self.build_vgg_forward(self.right,'resnet50-forward',reuse=True)
                else:
                    self.invDepth_left_est= self.build_resnet50_forward(self.left,'resnet50-forward',reuse=False)
                    if self.mode == 'train':
                        self.invDepth_right_est= self.build_resnet50_forward(self.right,'resnet50-forward',reuse=True)
//  [python]
def metric_fn(answers, prediction, start, end, yp1, yp2, num_answers):
  """Compute span accuracies and token F1/EM scores."""

  yp1 = tf.expand_dims(yp1, -1)
  yp2 = tf.expand_dims(yp2, -1)
  answer_mask = tf.sequence_mask(num_answers)

  start = tf.to_int64(start)
  end = tf.to_int64(end)
  start_correct = tf.reduce_any(tf.equal(start, yp1) & answer_mask, 1)
  end_correct = tf.reduce_any(tf.equal(end, yp2) & answer_mask, 1)
  correct = start_correct & end_correct

  em = tf.py_func(
      enum_fn(_exact_match_score, dtype='float32'),
      [prediction, answers, answer_mask], 'float32')
  f1 = tf.py_func(
      enum_fn(_f1_score, dtype='float32'), [prediction, answers, answer_mask],
      'float32')

  eval_metric_ops = {
      # TODO(ddohan): Add other useful metrics
      'acc_start':
          tf.metrics.mean(tf.cast(start_correct, 'float')),
      'acc_end':
          tf.metrics.mean(tf.cast(end_correct, 'float')),
      'acc_span':
          tf.metrics.mean(tf.cast(correct, 'float')),
      'em':
          tf.metrics.mean(em),
      'f1':
          tf.metrics.mean(f1),
      # Number of questions processed
      'num_question':
          tf.metrics.true_positives(
              tf.ones([tf.shape(prediction)][0]),
              tf.ones([tf.shape(prediction)][0]))
  }
  return eval_metric_ops
//  [python]
def metric_fn(answers, prediction, start, end, yp1, yp2, num_answers):
  """Compute span accuracies and token F1/EM scores."""

  yp1 = tf.expand_dims(yp1, -1)
  yp2 = tf.expand_dims(yp2, -1)
  answer_mask = tf.sequence_mask(num_answers)

  start = tf.cast(start, 'int64')
  end = tf.cast(end, 'int64')
  start_correct = tf.reduce_any(tf.equal(start, yp1) & answer_mask, 1)
  end_correct = tf.reduce_any(tf.equal(end, yp2) & answer_mask, 1)
  correct = start_correct & end_correct

  em = tf.py_func(
      enum_fn(_exact_match_score, dtype='float32'),
      [prediction, answers, answer_mask], 'float32')
  f1 = tf.py_func(
      enum_fn(_f1_score, dtype='float32'), [prediction, answers, answer_mask],
      'float32')

  eval_metric_ops = {
      # TODO(ddohan): Add other useful metrics
      'acc_start':
          tf.metrics.mean(tf.cast(start_correct, 'float')),
      'acc_end':
          tf.metrics.mean(tf.cast(end_correct, 'float')),
      'acc_span':
          tf.metrics.mean(tf.cast(correct, 'float')),
      'em':
          tf.metrics.mean(em),
      'f1':
          tf.metrics.mean(f1),
      # Number of questions processed
      'num_question':
          tf.metrics.true_positives(
              tf.ones([tf.shape(prediction)][0]),
              tf.ones([tf.shape(prediction)][0]))
  }
  return eval_metric_ops
//  [python]
def resample_example(example, max_length=256):
  """Given an example and max length, resample the context to that length.

  Start position randomly chosen from [0, answer_start]. Assumes a single
    answer per context, which is true for the SQuAD training set.

  Args:
    example: A single example containing at least these fields:
      ['answers_start_token', 'answers_end_token', 'context_tokens',
      'context_length']
    max_length: Maximum length. Contexts are resampled to this length.

  Returns:
    Resampled example.
  """

  # TODO(ddohan): Consider randomly cropping to shorter lengths
  # TODO(ddohan): Figure out how to resample the raw text as well. Not necessary
  # for training
  def _resample():
    """Helper method for resampling inside cond."""
    x = example
    ans_start = tf.to_int64(x['answers_start_token'][0])
    ans_end = tf.to_int64(x['answers_end_token'][0])
    min_start = tf.maximum(tf.to_int64(0), ans_end - max_length + 1)
    max_start = ans_start
    start_idx = tf.random_uniform([],
                                  min_start,
                                  max_start + 1, dtype=tf.int64)
    for k in ['answers_start_token', 'answers_end_token']:
      x[k] -= start_idx
    x['context_tokens'] = x['context_tokens'][start_idx:start_idx + max_length]
    x['context_length'] = tf.to_int64(tf.shape(x['context_tokens'])[0])
    return x

  def identity():
    return example

  return tf.cond(
      tf.greater_equal(
          tf.to_int32(max_length), tf.to_int32(example['context_length'])),
      true_fn=identity,
      false_fn=_resample)
//  [python]
def resample_example(example, max_length=256):
  """Given an example and max length, resample the context to that length.

  Start position randomly chosen from [0, answer_start]. Assumes a single
    answer per context, which is true for the SQuAD training set.

  Args:
    example: A single example containing at least these fields:
      ['answers_start_token', 'answers_end_token', 'context_tokens',
      'context_length']
    max_length: Maximum length. Contexts are resampled to this length.

  Returns:
    Resampled example.
  """

  # TODO(ddohan): Consider randomly cropping to shorter lengths
  # TODO(ddohan): Figure out how to resample the raw text as well. Not necessary
  # for training
  def _resample():
    """Helper method for resampling inside cond."""
    x = example
    ans_start = tf.cast(x['answers_start_token'][0], tf.int64)
    ans_end = tf.cast(x['answers_end_token'][0], tf.int64)
    min_start = tf.maximum(tf.cast(0, tf.int64), ans_end - max_length + 1)
    max_start = ans_start
    start_idx = tf.random_uniform([],
                                  min_start,
                                  max_start + 1, dtype=tf.int64)
    for k in ['answers_start_token', 'answers_end_token']:
      x[k] -= start_idx
    x['context_tokens'] = x['context_tokens'][start_idx:start_idx + max_length]
    x['context_length'] = tf.cast(tf.shape(x['context_tokens'])[0], tf.int64)
    return x

  def identity():
    return example

  return tf.cond(
      tf.greater_equal(
          tf.cast(max_length, tf.int32), tf.cast(example['context_length'], tf.int32)),
      true_fn=identity,
      false_fn=_resample)
//  [python]
def logp(self, x):
        return - 0.5 * U.sum(tf.square((x - self.mean) / self.std), axis=1) \
               - 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) \
               - U.sum(self.logstd, axis=1)
//  [python]
def logp(self, x):
        return - 0.5 * U.sum(tf.square((x - self.mean) / self.std), axis=1) \
               - 0.5 * np.log(2.0 * np.pi) * tf.cast(tf.shape(x)[1], tf.float32) \
               - U.sum(self.logstd, axis=1)
//  [python]
def logp(self, x):
        return - U.sum(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.logits, labels=tf.to_float(x)), axis=1)
//  [python]
def logp(self, x):
        return - U.sum(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.logits, labels=tf.cast(x, tf.float32)), axis=1)
//  [python]
def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["question_type"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["supporting_labels"] = tf.FixedLenFeature(
        [seq_length], tf.int64)
  name_to_features["supporting_mask"] = tf.FixedLenFeature(
      [seq_length], tf.int64)

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["question_type"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["supporting_labels"] = tf.FixedLenFeature(
        [seq_length], tf.int64)
  name_to_features["supporting_mask"] = tf.FixedLenFeature(
      [seq_length], tf.int64)

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def sample(self):
        p = tf.sigmoid(self.logits)
        u = tf.random_uniform(tf.shape(p))
        return tf.to_float(math_ops.less(u, p))
//  [python]
def sample(self):
        p = tf.sigmoid(self.logits)
        u = tf.random_uniform(tf.shape(p))
        return tf.cast(math_ops.less(u, p), tf.float32)
//  [python]
def preprocess(self, inputs):
    frames, actions = inputs
    frames = tf.image.convert_image_dtype(frames, tf.float32)
    if self.is_discrete_action:
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.to_float(actions)
    return frames, actions