jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_sequence_length(self, sequence):
        used = tf.sign(tf.reduce_max(tf.abs(sequence), axis=2))
        length = tf.reduce_sum(used, axis=1)
        length = tf.to_int32(length)

        return length
//  [python]
def compute_sequence_length(self, sequence):
        used = tf.sign(tf.reduce_max(tf.abs(sequence), axis=2))
        length = tf.reduce_sum(used, axis=1)
        length = tf.cast(length, tf.int32)

        return length
//  [python]
def _parse_eval_data(self, data):
    """Generates images and labels that are usable for model evaluation.

    Args:
      data: the decoded tensor dictionary from TfExampleDecoder.

    Returns:
      images: the image tensor.
      labels: a dict of Tensors that contains labels.
    """
    image = tf.cast(data['image'], dtype=tf.float32)
    boxes = data['groundtruth_boxes']
    classes = data['groundtruth_classes']

    image_shape = tf.shape(input=image)[0:2]
    # Converts boxes from normalized coordinates to pixel coordinates.
    boxes = box_ops.denormalize_boxes(boxes, image_shape)

    # Resizes and crops image.
    image, image_info = preprocess_ops.resize_and_crop_image(
        image,
        [self._output_height, self._output_width],
        padded_size=[self._output_height, self._output_width],
        aug_scale_min=1.0,
        aug_scale_max=1.0)
    unpad_image_shape = tf.cast(tf.shape(image), tf.float32)

    # Resizes and crops boxes.
    image_scale = image_info[2, :]
    offset = image_info[3, :]
    boxes = preprocess_ops.resize_and_crop_boxes(boxes, image_scale,
                                                 image_info[1, :], offset)

    # Filters out ground truth boxes that are all zeros.
    indices = box_ops.get_non_empty_box_indices(boxes)
    boxes = tf.gather(boxes, indices)
    classes = tf.gather(classes, indices)

    labels = self._build_label(
        unpad_image_shape=unpad_image_shape,
        boxes=boxes,
        classes=classes,
        image_info=image_info,
        data=data)

    if self._bgr_ordering:
      red, green, blue = tf.unstack(image, num=3, axis=2)
      image = tf.stack([blue, green, red], axis=2)

    image = preprocess_ops.normalize_image(
        image=image,
        offset=self._channel_means,
        scale=self._channel_stds)

    image = tf.cast(image, self._dtype)

    return image, labels
//  [python]
def intersection(gt_boxes, boxes):
  """Compute pairwise intersection areas between boxes.

  B: batch_size
  N: number of groundtruth boxes.
  M: number of anchor boxes.

  Args:
    gt_boxes: a float Tensor with [N, 4], or [B, N, 4]
    boxes: a float Tensor with [M, 4], or [B, M, 4]

  Returns:
    a float Tensor with shape [N, M] or [B, N, M] representing pairwise
      intersections.
  """
  with tf.name_scope('Intersection'):
    y_min1, x_min1, y_max1, x_max1 = tf.split(
        value=gt_boxes, num_or_size_splits=4, axis=-1)
    y_min2, x_min2, y_max2, x_max2 = tf.split(
        value=boxes, num_or_size_splits=4, axis=-1)

    boxes_rank = len(boxes.shape)
    perm = [1, 0] if boxes_rank == 2 else [0, 2, 1]
    # [N, M] or [B, N, M]
    y_min_max = tf.minimum(y_max1, tf.transpose(y_max2, perm))
    y_max_min = tf.maximum(y_min1, tf.transpose(y_min2, perm))
    x_min_max = tf.minimum(x_max1, tf.transpose(x_max2, perm))
    x_max_min = tf.maximum(x_min1, tf.transpose(x_min2, perm))

    intersect_heights = y_min_max - y_max_min
    intersect_widths = x_min_max - x_max_min
    zeros_t = tf.constant(0, dtype=intersect_heights.dtype)
    intersect_heights = tf.maximum(zeros_t, intersect_heights)
    intersect_widths = tf.maximum(zeros_t, intersect_widths)
    return intersect_heights * intersect_widths
//  [python]
def intersection(gt_boxes, boxes):
  """Compute pairwise intersection areas between boxes.

  B: batch_size
  N: number of groundtruth boxes.
  M: number of anchor boxes.

  Args:
    gt_boxes: a float Tensor with [N, 4], or [B, N, 4]
    boxes: a float Tensor with [M, 4], or [B, M, 4]

  Returns:
    a float Tensor with shape [N, M] or [B, N, M] representing pairwise
      intersections.
  """
  with tf.name_scope('Intersection'):
    y_min1, x_min1, y_max1, x_max1 = tf.split(
        value=gt_boxes, num_or_size_splits=4, axis=-1)
    y_min2, x_min2, y_max2, x_max2 = tf.split(
        value=boxes, num_or_size_splits=4, axis=-1)

    boxes_rank = len(boxes.shape)
    perm = [1, 0] if boxes_rank == 2 else [0, 2, 1]
    # [N, M] or [B, N, M]
    y_min_max = tf.minimum(y_max1, tf.transpose(y_max2, perm))
    y_max_min = tf.maximum(y_min1, tf.transpose(y_min2, perm))
    x_min_max = tf.minimum(x_max1, tf.transpose(x_max2, perm))
    x_max_min = tf.maximum(x_min1, tf.transpose(x_min2, perm))

    intersect_heights = y_min_max - y_max_min
    intersect_widths = x_min_max - x_max_min
    zeros_t = tf.cast(0, intersect_heights.dtype)
    intersect_heights = tf.maximum(zeros_t, intersect_heights)
    intersect_widths = tf.maximum(zeros_t, intersect_widths)
    return intersect_heights * intersect_widths
//  [python]
def _cross_suppression(boxes, box_slice, iou_threshold, inner_idx):
  batch_size = tf.shape(boxes)[0]
  new_slice = tf.slice(boxes, [0, inner_idx * NMS_TILE_SIZE, 0],
                       [batch_size, NMS_TILE_SIZE, 4])
  iou = box_ops.bbox_overlap(new_slice, box_slice)
  ret_slice = tf.expand_dims(
      tf.to_float(tf.reduce_all(iou < iou_threshold, [1])),
      2) * box_slice
  return boxes, ret_slice, iou_threshold, inner_idx + 1
//  [python]
def _cross_suppression(boxes, box_slice, iou_threshold, inner_idx):
  batch_size = tf.shape(boxes)[0]
  new_slice = tf.slice(boxes, [0, inner_idx * NMS_TILE_SIZE, 0],
                       [batch_size, NMS_TILE_SIZE, 4])
  iou = box_ops.bbox_overlap(new_slice, box_slice)
  ret_slice = tf.expand_dims(
      tf.cast(tf.reduce_all(iou < iou_threshold, [1]), box_slice.dtype),
      2) * box_slice
  return boxes, ret_slice, iou_threshold, inner_idx + 1
//  [python]
def deformable_conv2d(inputs, offset, ch_out, kernel_size=3,
                      activation=tf.identity, use_bias=True, W_init=None,
                      data_format="NHWC", scope="deformable"):
  with tf.variable_scope(scope):
    if data_format == "NCHW": # need NHWC
      inputs = tf.transpose(inputs, [0, 2, 3, 1])
      offset = tf.transpose(offset, [0, 2, 3, 1])
    in_shape = inputs.get_shape().as_list()
    in_channel = in_shape[3]
    input_h, input_w = tf.shape(inputs)[1], tf.shape(inputs)[2]

    shape = (kernel_size, kernel_size, in_channel, ch_out)
    kernel_n = shape[0] * shape[1]

    # [3x3, 2] local offset
    # (0, 0), (0, 1) ...
    # (1, 0), (1, 1) ..
    # ..
    initial_offset = tf.stack(
        tf.meshgrid(tf.range(shape[0]), tf.range(shape[1]), indexing="ij")
    )
    # [3x3, 2]
    initial_offset = tf.reshape(initial_offset, (-1, 2))
    # [1, 1, 3x3, 2]
    initial_offset = tf.expand_dims(tf.expand_dims(initial_offset, 0), 0)
    # [h, w, 3x3, 2]
    initial_offset = tf.tile(initial_offset, [input_h, input_w, 1, 1])
    initial_offset = t.cast(initial_offset, "float32")

    grid = tf.meshgrid(
        tf.range(-int((shape[0] - 1) / 2.0),
                 input_h - int((shape[0] - 1) / 2.0), 1),
        tf.range(-int((shape[0] - 1) / 2.0),
                 input_w - int((shape[0] - 1) / 2.0), 1),
        indexing="ij"
    )
    #[h, w, 2] # each h, w location?
    grid = tf.stack(grid, axis=-1)
    grid = t.cast(grid, "float32")
    grid = tf.expand_dims(grid, 2)
    # [h, w, 3*3, 2]
    grid = tf.tile(grid, [1, 1, kernel_n, 1])
    grid_offset = grid + initial_offset

    # [b, h, w, 3x3, c]
    input_deform = _tf_batch_map_offsets(inputs, offset, grid_offset)

    if W_init is None:
      W_init = tf.variance_scaling_initializer(scale=2.0)
    W = tf.get_variable(
        name="W", shape=[shape[0], shape[1], shape[-2], shape[-1]],
        initializer=W_init)

    W = tf.reshape(W, [1, 1, shape[0]*shape[1], shape[-2], shape[-1]])

    output = tf.nn.conv3d(
        input_deform, W, strides=[1, 2, 2, 1, 1], data_format="NDHWC",
        padding="VALID", name=None)
    # output is [b, new_h, new_w, 1, c]

    output = tf.squeeze(output, axis=3) # [b, h, w, c]
    if data_format == "NCHW":
      output = tf.transpose(output, [0, 3, 1, 2])

    if use_bias:
      b_init = tf.constant_initializer()
      b = tf.get_variable("b", [ch_out], initializer=b_init)
      output = tf.nn.bias_add(output, b, data_format=data_format)

    output = activation(output, name="output")
    return output
//  [python]
def deformable_conv2d(inputs, offset, ch_out, kernel_size=3,
                      activation=tf.identity, use_bias=True, W_init=None,
                      data_format="NHWC", scope="deformable"):
  with tf.variable_scope(scope):
    if data_format == "NCHW": # need NHWC
      inputs = tf.transpose(inputs, [0, 2, 3, 1])
      offset = tf.transpose(offset, [0, 2, 3, 1])
    in_shape = inputs.get_shape().as_list()
    in_channel = in_shape[3]
    input_h, input_w = tf.shape(inputs)[1], tf.shape(inputs)[2]

    shape = (kernel_size, kernel_size, in_channel, ch_out)
    kernel_n = shape[0] * shape[1]

    # [3x3, 2] local offset
    # (0, 0), (0, 1) ...
    # (1, 0), (1, 1) ..
    # ..
    initial_offset = tf.stack(
        tf.meshgrid(tf.range(shape[0]), tf.range(shape[1]), indexing="ij")
    )
    # [3x3, 2]
    initial_offset = tf.reshape(initial_offset, (-1, 2))
    # [1, 1, 3x3, 2]
    initial_offset = tf.expand_dims(tf.expand_dims(initial_offset, 0), 0)
    # [h, w, 3x3, 2]
    initial_offset = tf.tile(initial_offset, [input_h, input_w, 1, 1])
    initial_offset = tf.cast(initial_offset, "float32")

    grid = tf.meshgrid(
        tf.range(-int((shape[0] - 1) / 2.0),
                 input_h - int((shape[0] - 1) / 2.0), 1),
        tf.range(-int((shape[0] - 1) / 2.0),
                 input_w - int((shape[0] - 1) / 2.0), 1),
        indexing="ij"
    )
    #[h, w, 2] # each h, w location?
    grid = tf.stack(grid, axis=-1)
    grid = tf.cast(grid, "float32")
    grid = tf.expand_dims(grid, 2)
    # [h, w, 3*3, 2]
    grid = tf.tile(grid, [1, 1, kernel_n, 1])
    grid_offset = grid + initial_offset

    # [b, h, w, 3x3, c]
    input_deform = _tf_batch_map_offsets(inputs, offset, grid_offset)

    if W_init is None:
      W_init = tf.variance_scaling_initializer(scale=2.0)
    W = tf.get_variable(
        name="W", shape=[shape[0], shape[1], shape[-2], shape[-1]],
        initializer=W_init)

    W = tf.reshape(W, [1, 1, shape[0]*shape[1], shape[-2], shape[-1]])

    output = tf.nn.conv3d(
        input_deform, W, strides=[1, 2, 2, 1, 1], data_format="NDHWC",
        padding="VALID", name=None)
    # output is [b, new_h, new_w, 1, c]

    output = tf.squeeze(output, axis=3) # [b, h, w, c]
    if data_format == "NCHW":
      output = tf.transpose(output, [0, 3, 1, 2])

    if use_bias:
      b_init = tf.constant_initializer()
      b = tf.get_variable("b", [ch_out], initializer=b_init)
      output = tf.nn.bias_add(output, b, data_format=data_format)

    output = activation(output, name="output")
    return output
//  [python]
def __call__(self, boxes_1, boxes_2, boxes_1_masks=None, boxes_2_masks=None):
    """Compute pairwise IOU similarity between ground truth boxes and anchors.

    B: batch_size
    N: Number of groundtruth boxes.
    M: Number of anchor boxes.

    Args:
      boxes_1: a float Tensor with M or B * M boxes.
      boxes_2: a float Tensor with N or B * N boxes, the rank must be less than
        or equal to rank of `boxes_1`.
      boxes_1_masks: a boolean Tensor with M or B * M boxes. Optional.
      boxes_2_masks: a boolean Tensor with N or B * N boxes. Optional.

    Returns:
      A Tensor with shape [M, N] or [B, M, N] representing pairwise
        iou scores, anchor per row and groundtruth_box per colulmn.

    Input shape:
      boxes_1: [N, 4], or [B, N, 4]
      boxes_2: [M, 4], or [B, M, 4]
      boxes_1_masks: [N, 1], or [B, N, 1]
      boxes_2_masks: [M, 1], or [B, M, 1]

    Output shape:
      [M, N], or [B, M, N]
    """
    boxes_1 = t.cast(boxes_1, t.float32)
    boxes_2 = t.cast(boxes_2, t.float32)

    boxes_1_rank = len(boxes_1.shape)
    boxes_2_rank = len(boxes_2.shape)
    if boxes_1_rank < 2 or boxes_1_rank > 3:
      raise ValueError(
          '`groudtruth_boxes` must be rank 2 or 3, got {}'.format(boxes_1_rank))
    if boxes_2_rank < 2 or boxes_2_rank > 3:
      raise ValueError(
          '`anchors` must be rank 2 or 3, got {}'.format(boxes_2_rank))
    if boxes_1_rank < boxes_2_rank:
      raise ValueError('`groundtruth_boxes` is unbatched while `anchors` is '
                       'batched is not a valid use case, got groundtruth_box '
                       'rank {}, and anchors rank {}'.format(
                           boxes_1_rank, boxes_2_rank))

    result = iou(boxes_1, boxes_2)
    if boxes_1_masks is None and boxes_2_masks is None:
      return result
    background_mask = None
    mask_val_t = t.cast(self.mask_val, result.dtype) * t.ones_like(result)
    perm = [1, 0] if boxes_2_rank == 2 else [0, 2, 1]
    if boxes_1_masks is not None and boxes_2_masks is not None:
      background_mask = tf.logical_or(boxes_1_masks,
                                      tf.transpose(boxes_2_masks, perm))
    elif boxes_1_masks is not None:
      background_mask = boxes_1_masks
    else:
      background_mask = tf.logical_or(
          tf.zeros(tf.shape(boxes_2)[:-1], dtype=tf.bool),
          tf.transpose(boxes_2_masks, perm))
    return tf.where(background_mask, mask_val_t, result)
//  [python]
def __call__(self, boxes_1, boxes_2, boxes_1_masks=None, boxes_2_masks=None):
    """Compute pairwise IOU similarity between ground truth boxes and anchors.

    B: batch_size
    N: Number of groundtruth boxes.
    M: Number of anchor boxes.

    Args:
      boxes_1: a float Tensor with M or B * M boxes.
      boxes_2: a float Tensor with N or B * N boxes, the rank must be less than
        or equal to rank of `boxes_1`.
      boxes_1_masks: a boolean Tensor with M or B * M boxes. Optional.
      boxes_2_masks: a boolean Tensor with N or B * N boxes. Optional.

    Returns:
      A Tensor with shape [M, N] or [B, M, N] representing pairwise
        iou scores, anchor per row and groundtruth_box per colulmn.

    Input shape:
      boxes_1: [N, 4], or [B, N, 4]
      boxes_2: [M, 4], or [B, M, 4]
      boxes_1_masks: [N, 1], or [B, N, 1]
      boxes_2_masks: [M, 1], or [B, M, 1]

    Output shape:
      [M, N], or [B, M, N]
    """
    boxes_1 = tf.cast(boxes_1, tf.float32)
    boxes_2 = tf.cast(boxes_2, tf.float32)

    boxes_1_rank = len(boxes_1.shape)
    boxes_2_rank = len(boxes_2.shape)
    if boxes_1_rank < 2 or boxes_1_rank > 3:
      raise ValueError(
          '`groudtruth_boxes` must be rank 2 or 3, got {}'.format(boxes_1_rank))
    if boxes_2_rank < 2 or boxes_2_rank > 3:
      raise ValueError(
          '`anchors` must be rank 2 or 3, got {}'.format(boxes_2_rank))
    if boxes_1_rank < boxes_2_rank:
      raise ValueError('`groundtruth_boxes` is unbatched while `anchors` is '
                       'batched is not a valid use case, got groundtruth_box '
                       'rank {}, and anchors rank {}'.format(
                           boxes_1_rank, boxes_2_rank))

    result = iou(boxes_1, boxes_2)
    if boxes_1_masks is None and boxes_2_masks is None:
      return result
    background_mask = None
    mask_val_t = tf.cast(self.mask_val, result.dtype) * tf.ones_like(result)
    perm = [1, 0] if boxes_2_rank == 2 else [0, 2, 1]
    if boxes_1_masks is not None and boxes_2_masks is not None:
      background_mask = tf.logical_or(boxes_1_masks,
                                      tf.transpose(boxes_2_masks, perm))
    elif boxes_1_masks is not None:
      background_mask = boxes_1_masks
    else:
      background_mask = tf.logical_or(
          tf.zeros(tf.shape(boxes_2)[:-1], dtype=tf.bool),
          tf.transpose(boxes_2_masks, perm))
    return tf.where(background_mask, mask_val_t, result)
//  [python]
def preprocess(self, image: tf.Tensor,
                 label: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    """Apply image preprocessing and augmentation to the image and label."""
    if self.is_training:
      image = preprocessing.preprocess_for_train(
          image,
          image_size=self.image_size,
          mean_subtract=self.config.mean_subtract,
          standardize=self.config.standardize,
          dtype=self.dtype,
          augmenter=self.augmenter)
    else:
      image = preprocessing.preprocess_for_eval(
          image,
          image_size=self.image_size,
          num_channels=self.num_channels,
          mean_subtract=self.config.mean_subtract,
          standardize=self.config.standardize,
          dtype=self.dtype)

    label = tf.to_int32(label)
    if self.config.one_hot:
      label = tf.one_hot(label, self.num_classes)
      label = tf.reshape(label, [self.num_classes])

    return image, label
//  [python]
def preprocess(self, image: tf.Tensor,
                 label: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    """Apply image preprocessing and augmentation to the image and label."""
    if self.is_training:
      image = preprocessing.preprocess_for_train(
          image,
          image_size=self.image_size,
          mean_subtract=self.config.mean_subtract,
          standardize=self.config.standardize,
          dtype=self.dtype,
          augmenter=self.augmenter)
    else:
      image = preprocessing.preprocess_for_eval(
          image,
          image_size=self.image_size,
          num_channels=self.num_channels,
          mean_subtract=self.config.mean_subtract,
          standardize=self.config.standardize,
          dtype=self.dtype)

    label = tf.cast(label, tf.int32)
    if self.config.one_hot:
      label = tf.one_hot(label, self.num_classes)
      label = tf.reshape(label, [self.num_classes])

    return image, label
//  [python]
def inv_rot_weak_proj(p_orig, rot_orig,
                      rot):
  p_orig, rot_orig, rot = [
      t.cast(x, t.float64) for x in [p_orig, rot_orig, rot]]

  p_new = tf.einsum('jd,do->jo', p_orig, rot)
  rot_new = tf.einsum('kba,bc->kca', rot_orig, rot)
  return p_new, rot_new
//  [python]
def inv_rot_weak_proj(p_orig, rot_orig,
                      rot):
  p_orig, rot_orig, rot = [
      tf.cast(x, tf.float64) for x in [p_orig, rot_orig, rot]]

  p_new = tf.einsum('jd,do->jo', p_orig, rot)
  rot_new = tf.einsum('kba,bc->kca', rot_orig, rot)
  return p_new, rot_new
//  [python]
def inv_shift_weak_proj(p_orig, rot_orig,
                        s_orig, t_orig,
                        shift):
  p_orig, rot_orig, s_orig, t_orig, shift = [
      t.cast(x, t.float64) for x in [p_orig, rot_orig, s_orig, t_orig, shift]]

  p_new = p_orig - shift[None, :]
  new_sh = tf.einsum('jd,kdo->kjo', shift[None, :], rot_orig)[:, 0, :]
  t_new = t_orig + new_sh * s_orig[:, None]
  return p_new, t_new
//  [python]
def inv_shift_weak_proj(p_orig, rot_orig,
                        s_orig, t_orig,
                        shift):
  p_orig, rot_orig, s_orig, t_orig, shift = [
      tf.cast(x, tf.float64) for x in [p_orig, rot_orig, s_orig, t_orig, shift]]

  p_new = p_orig - shift[None, :]
  new_sh = tf.einsum('jd,kdo->kjo', shift[None, :], rot_orig)[:, 0, :]
  t_new = t_orig + new_sh * s_orig[:, None]
  return p_new, t_new
//  [python]
def inv_scale_weak_proj(p_orig, s_orig,
                        rescale):
  p_orig, s_orig, rescale = [
      t.cast(x, t.float64) for x in [p_orig, s_orig, rescale]]

  return p_orig / rescale, s_orig * rescale
//  [python]
def inv_scale_weak_proj(p_orig, s_orig,
                        rescale):
  p_orig, s_orig, rescale = [
      tf.cast(x, tf.float64) for x in [p_orig, s_orig, rescale]]

  return p_orig / rescale, s_orig * rescale
//  [python]
def call(self, inputs):
    x_inits, task_vecs = inputs['x_init'], inputs['task_vec']
    k_steps = len(self.models)
    batch_size = tf.shape(x_inits)[0]
    xs_cur = x_inits
    path = [xs_cur]
    for step_i in range(k_steps):
      fwds_cur = self.task_fwd_fn(xs_cur, task_vecs)
      vals_cur = self.task_eval_fn(fwds_cur, task_vecs)
      vals_cur = tf.reshape(vals_cur, (batch_size, 1))
      step_i_tensor = tf.convert_to_tensor(float(step_i))
      step_i_inputs = tf.repeat(step_i_tensor[None, None], batch_size, axis=0)
      inputs = [xs_cur, task_vecs, fwds_cur, vals_cur, step_i_inputs]
      inputs = [tf.to_float(x) for x in inputs]
      inputs_tensor = tf.concat(inputs, axis=-1)
      inputs_tensor = tf.stop_gradient(inputs_tensor)
      xs_update = self.models[step_i](inputs_tensor)
      xs_cur = tf.stop_gradient(xs_cur) + xs_update
      path.append(xs_cur)

    fwds_cur = self.task_fwd_fn(xs_cur, task_vecs)
    vals_cur = self.task_eval_fn(fwds_cur, task_vecs)
    last_task_loss = tf.reduce_mean(vals_cur)

    self.add_named_loss(last_task_loss, name=self.task_loss_name)

    outputs = {'x_opt': xs_cur, 'fwd': fwds_cur}

    for output_name, custom_fn in self.custom_output_fns.items():
      outputs[output_name] = custom_fn(xs_cur, task_vecs)

    return outputs
//  [python]
def call(self, inputs):
    x_inits, task_vecs = inputs['x_init'], inputs['task_vec']
    k_steps = len(self.models)
    batch_size = tf.shape(x_inits)[0]
    xs_cur = x_inits
    path = [xs_cur]
    for step_i in range(k_steps):
      fwds_cur = self.task_fwd_fn(xs_cur, task_vecs)
      vals_cur = self.task_eval_fn(fwds_cur, task_vecs)
      vals_cur = tf.reshape(vals_cur, (batch_size, 1))
      step_i_tensor = tf.convert_to_tensor(float(step_i))
      step_i_inputs = tf.repeat(step_i_tensor[None, None], batch_size, axis=0)
      inputs = [xs_cur, task_vecs, fwds_cur, vals_cur, step_i_inputs]
      inputs = [tf.cast(x, tf.float64) for x in inputs]
      inputs_tensor = tf.concat(inputs, axis=-1)
      inputs_tensor = tf.stop_gradient(inputs_tensor)
      xs_update = self.models[step_i](inputs_tensor)
      xs_cur = tf.stop_gradient(xs_cur) + xs_update
      path.append(xs_cur)

    fwds_cur = self.task_fwd_fn(xs_cur, task_vecs)
    vals_cur = self.task_eval_fn(fwds_cur, task_vecs)
    last_task_loss = tf.reduce_mean(vals_cur)

    self.add_named_loss(last_task_loss, name=self.task_loss_name)

    outputs = {'x_opt': xs_cur, 'fwd': fwds_cur}

    for output_name, custom_fn in self.custom_output_fns.items():
      outputs[output_name] = custom_fn(xs_cur, task_vecs)

    return outputs
//  [python]
def get_percentage_of_correct_keypoints(
    shapes_specs,
    threshold = 150,
    allow_mirror = True):
  """Vectorized PCK metric."""
  x_spec = shapes_specs[0]
  def pck(gt_pose_cam0, cur_x):
    # gt_pose_cam0: (17, 3), cur_x: (x_shape,)
    x_pose, x_rot = unpack_and_unparam(cur_x, x_spec)[:2]
    cam0_rot_pred = x_pose @ x_rot[0]
    errs = single_npjpe_not_mean(gt_pose_cam0, cam0_rot_pred, allow_mirror)
    pck = tf.reduce_mean(tf.to_float(errs < threshold)) * 100
    return pck

  pck_batch = vmap.vmap(pck, ['gt_pose_cam0', 'cur_x'])
  return pck_batch
//  [python]
def get_percentage_of_correct_keypoints(
    shapes_specs,
    threshold = 150,
    allow_mirror = True):
  """Vectorized PCK metric."""
  x_spec = shapes_specs[0]
  def pck(gt_pose_cam0, cur_x):
    # gt_pose_cam0: (17, 3), cur_x: (x_shape,)
    x_pose, x_rot = unpack_and_unparam(cur_x, x_spec)[:2]
    cam0_rot_pred = x_pose @ x_rot[0]
    errs = single_npjpe_not_mean(gt_pose_cam0, cam0_rot_pred, allow_mirror)
    pck = tf.reduce_mean(tf.cast(errs < threshold, tf.float32)) * 100
    return pck

  pck_batch = vmap.vmap(pck, ['gt_pose_cam0', 'cur_x'])
  return pck_batch
//  [python]
def _build_losses(self, inputs, model_out, compute_mask=False):
        images, labels = inputs
        regression_layer = self._config_dict['regression_layer']
        ss = 2 ** regression_layer

        gt_cls = labels['class']
        pred_cls = model_out['cls']

        gt_offsets = tf.nn.avg_pool2d(labels['dist_map'], ss, ss, 'SAME') / ss
        gt_sizes = tf.nn.avg_pool2d(labels['size_map'], ss, ss, 'SAME') / ss
        gt_weights = tf.nn.avg_pool2d(labels['weights'], ss, ss, 'SAME')

        offsets = model_out['offsets']
        sizes = model_out['sizes']
        weights = model_out['weights']

        w = tf.to_float(gt_weights[...,0] > 0.1)
        ofs_loss = tf.losses.huber(gt_offsets, offsets, 200)
        ofs_loss = tf.reduce_mean(tf.reduce_sum(ofs_loss * w, axis=(1,2))/tf.reduce_sum(w, axis=(1,2)))*ss*ss

        w = tf.to_float(gt_weights[...,0] > 0.5)
        size_loss = tf.losses.huber(gt_sizes, sizes, 200)
        size_loss = 0.25 * ss * ss * tf.reduce_mean(tf.reduce_sum(size_loss * w, axis=(1,2)) / tf.reduce_sum(w, axis=(1,2)))

        w = tf.to_float(gt_weights[...,0] >= 0)
        weight_loss = tf.losses.binary_crossentropy(gt_weights, weights, from_logits=True)
        weight_loss = tf.reduce_mean(tf.reduce_sum(weight_loss * w, axis=(1,2)) / tf.reduce_sum(w, axis=(1,2)))

        classification_loss = tf.reduce_mean(tf.losses.categorical_crossentropy(tf.one_hot(gt_cls, self._config_dict['n_cls']), pred_cls))

        model_loss = ofs_loss + weight_loss + size_loss + classification_loss
        #model_loss = ofs_loss + weight_loss + size_loss

        losses = {
            'ofs_loss': ofs_loss,
            'size_loss': size_loss,
            'weight_loss': weight_loss,
            'classification_loss': classification_loss,
        }
        if compute_mask:
            proposal_losses = self._proposal_layer.build_losses(labels, model_out)
            model_loss += proposal_losses['proposal_score_loss']
            model_loss += proposal_losses['proposal_loss']
            losses.update(proposal_losses)
            if self._config_dict['with_mask']:
                mask_loss = self._mask_layer._build_losses(labels, model_out)
                model_loss += mask_loss
                losses.update({
                    'mask_loss': mask_loss,
                })
        losses.update({
            'model_loss': model_loss,
        })
        return losses
//  [python]
def _build_losses(self, inputs, model_out, compute_mask=False):
        images, labels = inputs
        regression_layer = self._config_dict['regression_layer']
        ss = 2 ** regression_layer

        gt_cls = labels['class']
        pred_cls = model_out['cls']

        gt_offsets = tf.nn.avg_pool2d(labels['dist_map'], ss, ss, 'SAME') / ss
        gt_sizes = tf.nn.avg_pool2d(labels['size_map'], ss, ss, 'SAME') / ss
        gt_weights = tf.nn.avg_pool2d(labels['weights'], ss, ss, 'SAME')

        offsets = model_out['offsets']
        sizes = model_out['sizes']
        weights = model_out['weights']

        w = tf.cast(gt_weights[...,0] > 0.1, tf.float32)
        ofs_loss = tf.losses.huber(gt_offsets, offsets, 200)
        ofs_loss = tf.reduce_mean(tf.reduce_sum(ofs_loss * w, axis=(1,2))/tf.reduce_sum(w, axis=(1,2)))*ss*ss

        w = tf.cast(gt_weights[...,0] > 0.5, tf.float32)
        size_loss = tf.losses.huber(gt_sizes, sizes, 200)
        size_loss = 0.25 * ss * ss * tf.reduce_mean(tf.reduce_sum(size_loss * w, axis=(1,2)) / tf.reduce_sum(w, axis=(1,2)))

        w = tf.cast(gt_weights[...,0] >= 0, tf.float32)
        weight_loss = tf.losses.binary_crossentropy(gt_weights, weights, from_logits=True)
        weight_loss = tf.reduce_mean(tf.reduce_sum(weight_loss * w, axis=(1,2)) / tf.reduce_sum(w, axis=(1,2)))

        classification_loss = tf.reduce_mean(tf.losses.categorical_crossentropy(tf.one_hot(gt_cls, self._config_dict['n_cls']), pred_cls))

        model_loss = ofs_loss + weight_loss + size_loss + classification_loss
        #model_loss = ofs_loss + weight_loss + size_loss

        losses = {
            'ofs_loss': ofs_loss,
            'size_loss': size_loss,
            'weight_loss': weight_loss,
            'classification_loss': classification_loss,
        }
        if compute_mask:
            proposal_losses = self._proposal_layer.build_losses(labels, model_out)
            model_loss += proposal_losses['proposal_score_loss']
            model_loss += proposal_losses['proposal_loss']
            losses.update(proposal_losses)
            if self._config_dict['with_mask']:
                mask_loss = self._mask_layer._build_losses(labels, model_out)
                model_loss += mask_loss
                losses.update({
                    'mask_loss': mask_loss,
                })
        losses.update({
            'model_loss': model_loss,
        })
        return losses
//  [python]
def _cast_objects_class(inputs):
  """Set the object rotation and semantic properties.

  Args:
    inputs: A dictionary containing input tensors.
  """
  if standard_fields.InputDataFields.objects_class in inputs:
    inputs[standard_fields.InputDataFields.objects_class] = tf.to_int32(
        inputs[standard_fields.InputDataFields.objects_class])
//  [python]
def _cast_objects_class(inputs):
  """Set the object rotation and semantic properties.

  Args:
    inputs: A dictionary containing input tensors.
  """
  if standard_fields.InputDataFields.objects_class in inputs:
    inputs[standard_fields.InputDataFields.objects_class] = tf.cast(
        inputs[standard_fields.InputDataFields.objects_class], dtype=tf.int32)
//  [python]
def print_single_view_stats(opt_stats_ds):
  """Prints pose estimation metrics of monocular 3D estimates."""
  def _map_pmpjpe_nmpjpe_pck(key, rec):
    del key
    output = []
    for i in range(_N_CAM.value):
      epi_cam_f32 = t.cast(rec['pose3d_epi_pred'][i], t.float64)
      rot = t.cast(t.transpose(rec['cam_rot'][i]), t.float64)
      npje = single_npjpe_not_mean(
          rec['pose3d'] @ rot, epi_cam_f32, _DEBUG_ALLOW_MIRROR_IN_NMPJES.value)

      errs = [
          single_pmpjpe(rec['pose3d'], epi_cam_f32),
          t.reduce_mean(npje),
          t.reduce_mean(t.cast(npje < 150, t.float64)) * 100
      ]
      output.append(errs)
    per_cam_scores = [t.reduce_mean(x) for x in zip(*output)]
    return per_cam_scores

  init_scores_ds = opt_stats_ds.map(_map_pmpjpe_nmpjpe_pck)
  score = t.reduce_mean(list(init_scores_ds), axis=0)
  print('per-frame init pmpje/npjpe/pck: %.3f %.3f %.3f' % tuple(score.numpy()))
//  [python]
def print_single_view_stats(opt_stats_ds):
  """Prints pose estimation metrics of monocular 3D estimates."""
  def _map_pmpjpe_nmpjpe_pck(key, rec):
    del key
    output = []
    for i in range(_N_CAM.value):
      epi_cam_f32 = tf.cast(rec['pose3d_epi_pred'][i], tf.float64)
      rot = tf.cast(tf.transpose(rec['cam_rot'][i]), tf.float64)
      npje = single_npjpe_not_mean(
          rec['pose3d'] @ rot, epi_cam_f32, _DEBUG_ALLOW_MIRROR_IN_NMPJES.value)

      errs = [
          single_pmpjpe(rec['pose3d'], epi_cam_f32),
          tf.reduce_mean(npje),
          tf.reduce_mean(tf.cast(npje < 150, tf.float64)) * 100
      ]
      output.append(errs)
    per_cam_scores = [tf.reduce_mean(x) for x in zip(*output)]
    return per_cam_scores

  init_scores_ds = opt_stats_ds.map(_map_pmpjpe_nmpjpe_pck)
  score = tf.reduce_mean(list(init_scores_ds), axis=0)
  print('per-frame init pmpje/npjpe/pck: %.3f %.3f %.3f' % tuple(score.numpy()))
//  [python]
def _filter_valid_objects(inputs):
  """Removes the objects that do not contain 3d info.

  Args:
    inputs: A dictionary containing input tensors.
  """
  if standard_fields.InputDataFields.objects_class not in inputs:
    return

  valid_objects_mask = tf.reshape(
      tf.greater(inputs[standard_fields.InputDataFields.objects_class], 0),
      [-1])
  if standard_fields.InputDataFields.objects_has_3d_info in inputs:
    objects_with_3d_info = tf.reshape(
        tf.to_bool(
            inputs[standard_fields.InputDataFields.objects_has_3d_info]), [-1])
    valid_objects_mask = tf.logical_and(objects_with_3d_info,
                                        valid_objects_mask)
  if standard_fields.InputDataFields.objects_difficulty in inputs:
    valid_objects_mask = tf.logical_and(
        valid_objects_mask,
        tf.greater(
            tf.reshape(
                inputs[standard_fields.InputDataFields.objects_difficulty],
                [-1]), 0))
  for key in _OBJECT_KEYS:
    if key in inputs:
      inputs[key] = tf.boolean_mask(inputs[key], valid_objects_mask)
//  [python]
def _filter_valid_objects(inputs):
  """Removes the objects that do not contain 3d info.

  Args:
    inputs: A dictionary containing input tensors.
  """
  if standard_fields.InputDataFields.objects_class not in inputs:
    return

  valid_objects_mask = tf.reshape(
      tf.greater(inputs[standard_fields.InputDataFields.objects_class], 0),
      [-1])
  if standard_fields.InputDataFields.objects_has_3d_info in inputs:
    objects_with_3d_info = tf.reshape(
        tf.cast(
            inputs[standard_fields.InputDataFields.objects_has_3d_info],
            dtype=tf.bool), [-1])
    valid_objects_mask = tf.logical_and(objects_with_3d_info,
                                        valid_objects_mask)
  if standard_fields.InputDataFields.objects_difficulty in inputs:
    valid_objects_mask = tf.logical_and(
        valid_objects_mask,
        tf.greater(
            tf.reshape(
                inputs[standard_fields.InputDataFields.objects_difficulty],
                [-1]), 0))
  for key in _OBJECT_KEYS:
    if key in inputs:
      inputs[key] = tf.boolean_mask(inputs[key], valid_objects_mask)
//  [python]
def _adjust_class(self, tensor_dict: Mapping[str, Any]) -> tf.Tensor:
    """Change predicted class to what defiend by label.proto."""
    original_type = t.cast(tensor_dict['classes'], tf.uint8)
    if self._classes == 'all':
      adjusted_type = tf.where(
          tf.equal(original_type, 3),
          tf.ones_like(original_type) * 4,
          original_type)
    else:
      adjusted_type = tf.where(
          tf.equal(original_type, 1),
          tf.ones_like(original_type) * utils.CLASSES[self._classes],
          original_type)
    return adjusted_type
//  [python]
def _adjust_class(self, tensor_dict: Mapping[str, Any]) -> tf.Tensor:
    """Change predicted class to what defiend by label.proto."""
    original_type = tf.cast(tensor_dict['classes'], tf.uint8)
    if self._classes == 'all':
      adjusted_type = tf.where(
          tf.equal(original_type, 3),
          tf.ones_like(original_type) * 4,
          original_type)
    else:
      adjusted_type = tf.where(
          tf.equal(original_type, 1),
          tf.ones_like(original_type) * utils.CLASSES[self._classes],
          original_type)
    return adjusted_type
//  [python]
def call(self, t, states):
        logp_z = states[1]

        with tf.GradientTape(persistent=True) as inner_tape:
            z = states[0]
            inner_tape.watch(z)

            batchsize = z.shape[0]
            W, B, U = self.hyper_net(t)

            z = t.cast(z, W.dtype)
            Z = tf.tile(tf.expand_dims(z, axis=0), (self.width, 1, 1))

            h = tf.tanh(tf.matmul(Z, W) + B)
            dz_dt = tf.reduce_mean(tf.matmul(h, U), axis=0)
            dz_dt_sum = tf.reduce_sum(dz_dt, axis=0)

            outputs = [dz_dt_sum[i] for i in range(dz_dt_sum.shape[0])]

        """ Inlining to satisfy gradient tape
        Calculates the trace of the Jacobian df/dz.
        Obtained from: https://github.com/rtqichen/ffjord/blob/master/lib/layers/odefunc.py#L13
        """
        sum_diag = 0.
        for i in range(z.shape[1]):
            sum_diag += inner_tape.gradient(outputs[i], z)[:, i]

        dlogp_z_dt = -tf.reshape(sum_diag, (batchsize, 1))
        return (dz_dt, dlogp_z_dt)
//  [python]
def call(self, t, states):
        logp_z = states[1]

        with tf.GradientTape(persistent=True) as inner_tape:
            z = states[0]
            inner_tape.watch(z)

            batchsize = z.shape[0]
            W, B, U = self.hyper_net(t)

            z = tf.cast(z, W.dtype)
            Z = tf.tile(tf.expand_dims(z, axis=0), (self.width, 1, 1))

            h = tf.tanh(tf.matmul(Z, W) + B)
            dz_dt = tf.reduce_mean(tf.matmul(h, U), axis=0)
            dz_dt_sum = tf.reduce_sum(dz_dt, axis=0)

            outputs = [dz_dt_sum[i] for i in range(dz_dt_sum.shape[0])]

        """ Inlining to satisfy gradient tape
        Calculates the trace of the Jacobian df/dz.
        Obtained from: https://github.com/rtqichen/ffjord/blob/master/lib/layers/odefunc.py#L13
        """
        sum_diag = 0.
        for i in range(z.shape[1]):
            sum_diag += inner_tape.gradient(outputs[i], z)[:, i]

        dlogp_z_dt = -tf.reshape(sum_diag, (batchsize, 1))
        return (dz_dt, dlogp_z_dt)
//  [python]
def create_test_input(batch, height, width, channels):
  """Create test input tensor."""
  if None in [batch, height, width, channels]:
    return tf.placeholder(tf.float32, (batch, height, width, channels))
  else:
    return t.cast(
        np.tile(
            np.reshape(
                np.reshape(np.arange(height), [height, 1]) +
                np.reshape(np.arange(width), [1, width]),
                [1, height, width, 1]),
            [batch, 1, 1, channels]),
        tf.float32)
//  [python]
def create_test_input(batch, height, width, channels):
  """Create test input tensor."""
  if None in [batch, height, width, channels]:
    return tf.placeholder(tf.float32, (batch, height, width, channels))
  else:
    return tf.cast(
        np.tile(
            np.reshape(
                np.reshape(np.arange(height), [height, 1]) +
                np.reshape(np.arange(width), [1, width]),
                [1, height, width, 1]),
            [batch, 1, 1, channels]),
        tf.float32)
//  [python]
def test_rotate(self, dtype):
    image = tf.reshape(tf.cast(tf.range(9), dtype), (3, 3))
    rotation = 90.
    transformed = augment.rotate(image=image, degrees=rotation)
    expected = [[2, 5, 8], [1, 4, 7], [0, 3, 6]]
    self.assertAllEqual(transformed, expected)
//  [python]
def on_epoch_end(self, epoch, logs={}):
        num_thresholds = self.num_thresholds
        thresholds = [(i + 1) * 1.0 / (num_thresholds - 1) for i in range(num_thresholds - 2)]
        custom_val_metrics = list()

        y_pred = tf.reshape(self.model.predict(self.X_val), [-1])
        y_true = tf.reshape(self.y_val, [-1])

        for thresh in thresholds:
            pred_classes = tf.cast(tf.math.greater(y_pred, tf.constant([thresh])), dtype=tf.float32)
            cm = tf.math.confusion_matrix(y_true, pred_classes, num_classes=2)

            tn = cm[0][0].numpy()
            fp = cm[0][1].numpy()
            fn = cm[1][0].numpy()
            tp = cm[1][1].numpy()

            precision = tp / (tp + fp + K.epsilon())
            recall = tp / (tp + fn + K.epsilon())
            f1 = 2 * (precision * recall) / (precision + recall + K.epsilon())

            payload = \
                OrderedDict([('val_auc', logs.get('val_auc')),
                             ('val_f1', f1),
                             ('precision', precision),
                             ('recall', recall),
                             ('val_tp', tp),
                             ('val_tn', tn),
                             ('val_fp', fp),
                             ('val_fn', fn),
                             ('thresh', thresh)])

            custom_val_metrics.append(payload)
        best_val_f1 = sorted(custom_val_metrics, key=lambda x: x['val_f1'], reverse=True)[0]
        pprint.pprint(best_val_f1)
//  [python]
def one_hot(x, name):
        return tf.multiply(
            x,
            tf.cast(tf.equal(x, tf.expand_dims(tf.reduce_max(x, 1), 1)), tf.float32),
            name=name
        )
//  [python]
def build_relative_position(query_size, key_size):
    """
    Build relative position according to the query and key

    We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
    \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
    P_k\\)

    Args:
        query_size (int): the length of query
        key_size (int): the length of key

    Return:
        `tf.Tensor`: A tensor with shape [1, query_size, key_size]

    """
    q_ids = tf.range(query_size, dtype=tf.int32)
    k_ids = tf.range(key_size, dtype=tf.int32)
    rel_pos_ids = q_ids[:, None] - tf.tile(tf.reshape(k_ids, [1, -1]), [query_size, 1])
    rel_pos_ids = rel_pos_ids[:query_size, :]
    rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)
    return t.cast(rel_pos_ids, tf.int64)
//  [python]
def build_relative_position(query_size, key_size):
    """
    Build relative position according to the query and key

    We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
    \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
    P_k\\)

    Args:
        query_size (int): the length of query
        key_size (int): the length of key

    Return:
        `tf.Tensor`: A tensor with shape [1, query_size, key_size]

    """
    q_ids = tf.range(query_size, dtype=tf.int32)
    k_ids = tf.range(key_size, dtype=tf.int32)
    rel_pos_ids = q_ids[:, None] - tf.tile(tf.reshape(k_ids, [1, -1]), [query_size, 1])
    rel_pos_ids = rel_pos_ids[:query_size, :]
    rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)
    return tf.cast(rel_pos_ids, tf.int64)
//  [python]
def test_self_supervised_augmented_loss_identity(self):
    model = build_meta_arch(predict_full_resolution_masks=True,
                            augmented_self_supervision_max_translation=0.0)

    x = tf.random.uniform((2, 3, 32, 32), 0, 1)
    boxes = tf.constant([[0., 0., 1., 1.]] * 6)
    boxes = tf.reshape(boxes, [2, 3, 4])
    x = tf.cast(x > 0, tf.float32)
    x = (x - 0.5) * 2e40  # x is a tensor or large +ve or -ve values.
    loss = model._compute_self_supervised_augmented_loss(x, x, boxes)

    self.assertAlmostEqual(loss.numpy().sum(), 0.0)
//  [python]
def testSeparableConv2DSameWithInputEvenSize(self):
    n, n2 = 4, 2

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    dw = create_test_input(1, 3, 3, 1)
    dw = tf.reshape(dw, [3, 3, 1, 1])

    tf.get_variable('Conv/depthwise_weights', initializer=dw)
    tf.get_variable('Conv/pointwise_weights',
                    initializer=tf.ones([1, 1, 1, 1]))
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=1, scope='Conv')
    y1_expected = tf.constant([[14, 28, 43, 26],
                           [28, 48, 66, 37],
                           [43, 66, 84, 46],
                           [26, 37, 46, 22]], dtype=tf.float32)
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.constant([[14, 43],
                           [43, 84]], dtype=tf.float32)
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = xception.separable_conv2d_same(x, 1, 3, depth_multiplier=1,
                                        regularize_depthwise=True,
                                        stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=2, scope='Conv')
    y4_expected = tf.constant([[48, 37],
                           [37, 22]], dtype=tf.float32)
    y4_expected = tf.reshape(y4_expected, [1, n2, n2, 1])

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())
//  [python]
def testSeparableConv2DSameWithInputEvenSize(self):
    n, n2 = 4, 2

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    dw = create_test_input(1, 3, 3, 1)
    dw = tf.reshape(dw, [3, 3, 1, 1])

    tf.get_variable('Conv/depthwise_weights', initializer=dw)
    tf.get_variable('Conv/pointwise_weights',
                    initializer=tf.ones([1, 1, 1, 1]))
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=1, scope='Conv')
    y1_expected = tf.cast([[14, 28, 43, 26],
                           [28, 48, 66, 37],
                           [43, 66, 84, 46],
                           [26, 37, 46, 22]], tf.float32)
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.cast([[14, 43],
                           [43, 84]], tf.float32)
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = xception.separable_conv2d_same(x, 1, 3, depth_multiplier=1,
                                        regularize_depthwise=True,
                                        stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=2, scope='Conv')
    y4_expected = tf.cast([[48, 37],
                           [37, 22]], tf.float32)
    y4_expected = tf.reshape(y4_expected, [1, n2, n2, 1])

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())
//  [python]
def _one_phase_preprocessing_fn(inputs):
  x_plus_one = _plus_one(inputs['x'])
  subtracted = tf.sparse.add(
      t.cast(inputs['y'], tf.float32), -analyzers.mean(x_plus_one))
  _ = analyzers.vocabulary(inputs['s'])
  return {'subtracted': subtracted}
//  [python]
def _one_phase_preprocessing_fn(inputs):
  x_plus_one = _plus_one(inputs['x'])
  subtracted = tf.sparse.add(
      tf.cast(inputs['y'], tf.float32), -analyzers.mean(x_plus_one))
  _ = analyzers.vocabulary(inputs['s'])
  return {'subtracted': subtracted}
//  [python]
def testSeparableConv2DSameWithInputOddSize(self):
    n, n2 = 5, 3

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    dw = create_test_input(1, 3, 3, 1)
    dw = tf.reshape(dw, [3, 3, 1, 1])

    tf.get_variable('Conv/depthwise_weights', initializer=dw)
    tf.get_variable('Conv/pointwise_weights',
                    initializer=tf.ones([1, 1, 1, 1]))
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=1, scope='Conv')
    y1_expected = tf.constant([[14, 28, 43, 58, 34],
                           [28, 48, 66, 84, 46],
                           [43, 66, 84, 102, 55],
                           [58, 84, 102, 120, 64],
                           [34, 46, 55, 64, 30]], dtype=tf.float32)
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.constant([[14, 43, 34],
                           [43, 84, 55],
                           [34, 55, 30]], dtype=tf.float32)
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = xception.separable_conv2d_same(x, 1, 3, depth_multiplier=1,
                                        regularize_depthwise=True,
                                        stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=2, scope='Conv')
    y4_expected = y2_expected

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())