jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def augment_bounding_boxes(
        self, bounding_boxes, transformations, raw_images=None, **kwargs
    ):
        if self.bounding_box_format is None:
            raise ValueError(
                "Please provide a `bounding_box_format` when augmenting "
                "bounding boxes with `JitteredResize()`."
            )
        if isinstance(bounding_boxes["boxes"], tf.RaggedTensor):
            bounding_boxes = bounding_box.to_dense(bounding_boxes)
        result = bounding_boxes.copy()
        image_scales = t.cast(
            transformations["image_scales"], self.compute_dtype
        )
        offsets = t.cast(transformations["offsets"], self.compute_dtype)

        bounding_boxes = bounding_box.convert_format(
            bounding_boxes,
            images=raw_images,
            source=self.bounding_box_format,
            target="yxyx",
        )

        # Adjusts box coordinates based on image_scale and offset.
        yxyx = bounding_boxes["boxes"]
        yxyx *= t.tile(image_scales, [1, 2])[..., tf.newaxis, :]
        yxyx -= t.tile(offsets, [1, 2])[..., tf.newaxis, :]

        result["boxes"] = yxyx
        result = bounding_box.clip_to_image(
            result,
            image_shape=self.target_size + (3,),
            bounding_box_format="yxyx",
        )
        result = bounding_box.convert_format(
            result,
            image_shape=self.target_size + (3,),
            source="yxyx",
            target=self.bounding_box_format,
        )
        return result
//  [python]
def augment_bounding_boxes(
        self, bounding_boxes, transformations, raw_images=None, **kwargs
    ):
        if self.bounding_box_format is None:
            raise ValueError(
                "Please provide a `bounding_box_format` when augmenting "
                "bounding boxes with `JitteredResize()`."
            )
        if isinstance(bounding_boxes["boxes"], tf.RaggedTensor):
            bounding_boxes = bounding_box.to_dense(bounding_boxes)
        result = bounding_boxes.copy()
        image_scales = tf.cast(
            transformations["image_scales"], self.compute_dtype
        )
        offsets = tf.cast(transformations["offsets"], self.compute_dtype)

        bounding_boxes = bounding_box.convert_format(
            bounding_boxes,
            images=raw_images,
            source=self.bounding_box_format,
            target="yxyx",
        )

        # Adjusts box coordinates based on image_scale and offset.
        yxyx = bounding_boxes["boxes"]
        yxyx *= tf.tile(image_scales, [1, 2])[..., tf.newaxis, :]
        yxyx -= tf.tile(offsets, [1, 2])[..., tf.newaxis, :]

        result["boxes"] = yxyx
        result = bounding_box.clip_to_image(
            result,
            image_shape=self.target_size + (3,),
            bounding_box_format="yxyx",
        )
        result = bounding_box.convert_format(
            result,
            image_shape=self.target_size + (3,),
            source="yxyx",
            target=self.bounding_box_format,
        )
        return result
//  [python]
def resize_and_crop_single_image(self, inputs, resize_method="bilinear"):
        image = inputs.get("images", None)
        scaled_size = inputs.get("scaled_sizes", None)
        offset = inputs.get("offsets", None)

        scaled_image = tf.image.resize(
            image, tf.to_int32(scaled_size), method=resize_method
        )
        scaled_image = scaled_image[
            offset[0] : offset[0] + self.crop_size[0],
            offset[1] : offset[1] + self.crop_size[1],
            :,
        ]
        scaled_image = tf.image.pad_to_bounding_box(
            scaled_image, 0, 0, self.target_size[0], self.target_size[1]
        )
        return scaled_image
//  [python]
def resize_and_crop_single_image(self, inputs, resize_method="bilinear"):
        image = inputs.get("images", None)
        scaled_size = inputs.get("scaled_sizes", None)
        offset = inputs.get("offsets", None)

        scaled_image = tf.image.resize(
            image, tf.cast(scaled_size, tf.int32), method=resize_method
        )
        scaled_image = scaled_image[
            offset[0] : offset[0] + self.crop_size[0],
            offset[1] : offset[1] + self.crop_size[1],
            :,
        ]
        scaled_image = tf.image.pad_to_bounding_box(
            scaled_image, 0, 0, self.target_size[0], self.target_size[1]
        )
        return scaled_image
//  [python]
def clipoutNeg(vec, threshold=1e-6):
    mask = t.cast(vec > threshold, t.float32)
    return mask * vec
//  [python]
def clipoutNeg(vec, threshold=1e-6):
    mask = tf.cast(vec > threshold, tf.float32)
    return mask * vec
//  [python]
def _get_num_pos_neg_samples(self, sorted_indices_tensor, sample_size):
    """Counts the number of positives and negatives numbers to be sampled.

    Args:
      sorted_indices_tensor: A sorted int32 tensor of shape [N] which contains
        the signed indices of the examples where the sign is based on the label
        value. The examples that cannot be sampled are set to 0. It samples
        at most sample_size*positive_fraction positive examples and remaining
        from negative examples.
      sample_size: Size of subsamples.

    Returns:
      A tuple containing the number of positive and negative labels in the
      subsample.
    """
    input_length = tf.shape(input=sorted_indices_tensor)[0]
    valid_positive_index = tf.greater(sorted_indices_tensor,
                                      tf.zeros(input_length, tf.int32))
    num_sampled_pos = tf.reduce_sum(
        input_tensor=t.cast(valid_positive_index, tf.int32))
    max_num_positive_samples = t.cast(
        t.cast(sample_size, tf.float32) * self._positive_fraction, tf.int32)
    num_positive_samples = tf.minimum(max_num_positive_samples, num_sampled_pos)
    num_negative_samples = tf.constant(sample_size,
                                       tf.int32) - num_positive_samples

    return num_positive_samples, num_negative_samples
//  [python]
def _get_num_pos_neg_samples(self, sorted_indices_tensor, sample_size):
    """Counts the number of positives and negatives numbers to be sampled.

    Args:
      sorted_indices_tensor: A sorted int32 tensor of shape [N] which contains
        the signed indices of the examples where the sign is based on the label
        value. The examples that cannot be sampled are set to 0. It samples
        at most sample_size*positive_fraction positive examples and remaining
        from negative examples.
      sample_size: Size of subsamples.

    Returns:
      A tuple containing the number of positive and negative labels in the
      subsample.
    """
    input_length = tf.shape(input=sorted_indices_tensor)[0]
    valid_positive_index = tf.greater(sorted_indices_tensor,
                                      tf.zeros(input_length, tf.int32))
    num_sampled_pos = tf.reduce_sum(
        input_tensor=tf.cast(valid_positive_index, tf.int32))
    max_num_positive_samples = tf.cast(
        tf.cast(sample_size, tf.float32) * self._positive_fraction, tf.int32)
    num_positive_samples = tf.minimum(max_num_positive_samples, num_sampled_pos)
    num_negative_samples = tf.constant(sample_size,
                                       tf.int32) - num_positive_samples

    return num_positive_samples, num_negative_samples
//  [python]
def vecdot(
    x1: Union[tf.Tensor, tf.Variable],
    x2: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: int = -1,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    dtype = ivy.as_native_dtype(ivy.promote_types(x1.dtype, x2.dtype))
    return tf.cast(tf.tensordot(x1, x2, axes=(axis, axis)), dtype)
//  [python]
def tensordot(
    x1: Union[tf.Tensor, tf.Variable],
    x2: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axes: Union[int, Tuple[List[int], List[int]]] = 2,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    dtype = ivy.as_native_dtype(ivy.promote_types(x1.dtype, x2.dtype))
    ret = tf.cast(tf.tensordot(x1, x2, axes=axes), dtype)
    return ret
//  [python]
def _get_values_from_start_and_end(self, input_tensor, num_start_samples,
                                     num_end_samples, total_num_samples):
    """slices num_start_samples and last num_end_samples from input_tensor.

    Args:
      input_tensor: An int32 tensor of shape [N] to be sliced.
      num_start_samples: Number of examples to be sliced from the beginning of
        the input tensor.
      num_end_samples: Number of examples to be sliced from the end of the input
        tensor.
      total_num_samples: Sum of is num_start_samples and num_end_samples. This
        should be a scalar.

    Returns:
      A tensor containing the first num_start_samples and last num_end_samples
      from input_tensor.

    """
    input_length = tf.shape(input=input_tensor)[0]
    start_positions = tf.less(tf.range(input_length), num_start_samples)
    end_positions = tf.greater_equal(
        tf.range(input_length), input_length - num_end_samples)
    selected_positions = tf.logical_or(start_positions, end_positions)
    selected_positions = tf.to_float(selected_positions)
    indexed_positions = tf.multiply(
        tf.cumsum(selected_positions), selected_positions)
    one_hot_selector = tf.one_hot(
        tf.cast(indexed_positions, tf.int32) - 1,
        total_num_samples,
        dtype=tf.float32)
    return tf.to_int32(
        tf.tensordot(
            tf.to_float(input_tensor), one_hot_selector, axes=[0, 0]))
//  [python]
def _get_values_from_start_and_end(self, input_tensor, num_start_samples,
                                     num_end_samples, total_num_samples):
    """slices num_start_samples and last num_end_samples from input_tensor.

    Args:
      input_tensor: An int32 tensor of shape [N] to be sliced.
      num_start_samples: Number of examples to be sliced from the beginning of
        the input tensor.
      num_end_samples: Number of examples to be sliced from the end of the input
        tensor.
      total_num_samples: Sum of is num_start_samples and num_end_samples. This
        should be a scalar.

    Returns:
      A tensor containing the first num_start_samples and last num_end_samples
      from input_tensor.

    """
    input_length = tf.shape(input=input_tensor)[0]
    start_positions = tf.less(tf.range(input_length), num_start_samples)
    end_positions = tf.greater_equal(
        tf.range(input_length), input_length - num_end_samples)
    selected_positions = tf.logical_or(start_positions, end_positions)
    selected_positions = tf.cast(selected_positions, tf.float32)
    indexed_positions = tf.multiply(
        tf.cumsum(selected_positions), selected_positions)
    one_hot_selector = tf.one_hot(
        tf.cast(indexed_positions, tf.int32) - 1,
        total_num_samples,
        dtype=tf.float32)
    return tf.cast(
        tf.tensordot(
            tf.cast(input_tensor, tf.float32), one_hot_selector, axes=[0, 0]),
        tf.int32)
//  [python]
def reweight(msa1hot, cutoff):
    with tf.name_scope('reweight'):
        id_min = tf.cast(tf.shape(msa1hot)[1], tf.float32) * cutoff
        id_mtx = tf.tensordot(msa1hot, msa1hot, [[1,2], [1,2]])
        id_mask = id_mtx > id_min
        w = 1.0/tf.reduce_sum(tf.cast(id_mask, dtype=tf.float32),-1)
    return w
//  [python]
def vector_norm(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: Optional[Union[int, Sequence[int]]] = None,
    keepdims: bool = False,
    ord: Union[int, float, Literal[inf, -inf]] = 2,
    dtype: Optional[DType] = None,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    if dtype and x.dtype != dtype:
        x = tf.cast(x, dtype)
    abs_x = tf.abs(x)
    if ord == 0:
        return tf.reduce_sum(tf.cast(x != 0, abs_x.dtype), axis=axis, keepdims=keepdims)
    elif ord == inf:
        return tf.reduce_max(abs_x, axis=axis, keepdims=keepdims)
    elif ord == -inf:
        return tf.reduce_min(abs_x, axis=axis, keepdims=keepdims)
    else:
        return tf.reduce_sum(abs_x**ord, axis=axis, keepdims=keepdims) ** (1.0 / ord)
//  [python]
def matrix_norm(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    ord: Union[int, float, Literal[inf, -inf, "fro", "nuc"]] = "fro",
    axis: Tuple[int, int] = (-2, -1),
    keepdims: bool = False,
    dtype: Optional[tf.DType] = None,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    if dtype is not None:
        x = ivy.astype(x, dtype).to_native()
    if ord == "nuc":
        x = tf.experimental.numpy.moveaxis(x, axis, (-2, -1))
        ret = tf.reduce_sum(
            tf.linalg.svd(x, compute_uv=False),
            axis=-1,
        )
    elif ord == -1:
        ret = tf.reduce_min(
            tf.reduce_sum(tf.abs(x), axis=axis[0], keepdims=True),
            axis=axis,
            keepdims=keepdims,
        )
    elif ord == -2:
        x = tf.experimental.numpy.moveaxis(x, axis, (-2, -1))
        ret = tf.reduce_min(
            tf.linalg.svd(x, compute_uv=False),
            axis=-1,
        )
    elif ord == float("-inf"):
        ret = tf.reduce_min(
            tf.reduce_sum(tf.abs(x), axis=axis[1], keepdims=True),
            axis=axis,
            keepdims=keepdims,
        )
    else:
        ret = tf.norm(x, ord=ord, axis=axis, keepdims=keepdims)
    ret = tf.cast(ret, ret.dtype.real_dtype)
    if keepdims and ord in [-2, "nuc"]:
        for dim in axis:
            ret = tf.expand_dims(ret, dim % tf.rank(x))
    return ret
//  [python]
def _positions_center_origin(height, width):
  """Returns image coordinates where the origin at the image center."""
  h = tf.range(0.0, height, 1)
  w = tf.range(0.0, width, 1)
  center_h = tf.to_float(height) / 2.0 - 0.5
  center_w = tf.to_float(width) / 2.0 - 0.5
  return tf.stack(tf.meshgrid(h - center_h, w - center_w, indexing='ij'), -1)
//  [python]
def _positions_center_origin(height, width):
  """Returns image coordinates where the origin at the image center."""
  h = tf.range(0.0, height, 1)
  w = tf.range(0.0, width, 1)
  center_h = tf.cast(height, tf.float32) / 2.0 - 0.5
  center_w = tf.cast(width, tf.float32) / 2.0 - 0.5
  return tf.stack(tf.meshgrid(h - center_h, w - center_w, indexing='ij'), -1)
//  [python]
def _center_crop(images, height, width):
  """Performs a center crop with the given heights and width."""
  # ensure height, width to be int
  height = t.cast(height, t.int32)
  width = t.cast(width, t.int32)
  # get current size
  images_shape = tf.shape(images)
  current_height = images_shape[-3]
  current_width = images_shape[-2]
  # compute required offset
  offset_height = t.cast((current_height - height) / 2, t.int32)
  offset_width = t.cast((current_width - width) / 2, t.int32)
  # perform the crop
  images = tf.image.crop_to_bounding_box(
      images, offset_height, offset_width, height, width)
  return images
//  [python]
def _center_crop(images, height, width):
  """Performs a center crop with the given heights and width."""
  # ensure height, width to be int
  height = tf.cast(height, tf.int32)
  width = tf.cast(width, tf.int32)
  # get current size
  images_shape = tf.shape(images)
  current_height = images_shape[-3]
  current_width = images_shape[-2]
  # compute required offset
  offset_height = tf.cast((current_height - height) / 2, tf.int32)
  offset_width = tf.cast((current_width - width) / 2, tf.int32)
  # perform the crop
  images = tf.image.crop_to_bounding_box(
      images, offset_height, offset_width, height, width)
  return images
//  [python]
def random_scale(images, flow=None, mask=None, min_scale=1.0, max_scale=1.0):
  """Performs a random scaling in the given range."""
  # choose a random scale factor and compute new resolution
  orig_height = tf.shape(images)[-3]
  orig_width = tf.shape(images)[-2]
  scale = tf.random.uniform([],
                            minval=min_scale,
                            maxval=max_scale,
                            dtype=tf.float32)
  new_height = tf.to_int32(
      tf.math.ceil(tf.to_float(orig_height) * scale))
  new_width = tf.to_int32(
      tf.math.ceil(tf.to_float(orig_width) * scale))

  # rescale the images (and flow)
  images = uflow_utils.resize(images, new_height, new_width, is_flow=False)
  if flow is not None:
    flow, mask = uflow_utils.resize(
        flow, new_height, new_width, is_flow=True, mask=mask)
  return images, flow, mask
//  [python]
def random_scale(images, flow=None, mask=None, min_scale=1.0, max_scale=1.0):
  """Performs a random scaling in the given range."""
  # choose a random scale factor and compute new resolution
  orig_height = tf.shape(images)[-3]
  orig_width = tf.shape(images)[-2]
  scale = tf.random.uniform([],
                            minval=min_scale,
                            maxval=max_scale,
                            dtype=tf.float32)
  new_height = tf.cast(
      tf.math.ceil(tf.cast(orig_height, tf.float32) * scale), tf.int32)
  new_width = tf.cast(
      tf.math.ceil(tf.cast(orig_width, tf.float32) * scale), tf.int32)

  # rescale the images (and flow)
  images = uflow_utils.resize(images, new_height, new_width, is_flow=False)
  if flow is not None:
    flow, mask = uflow_utils.resize(
        flow, new_height, new_width, is_flow=True, mask=mask)
  return images, flow, mask
//  [python]
def matmul(
    x1: Union[tf.Tensor, tf.Variable],
    x2: Union[tf.Tensor, tf.Variable],
    /,
    *,
    transpose_a: bool = False,
    transpose_b: bool = False,
    adjoint_a: bool = False,
    adjoint_b: bool = False,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    x1, x2 = ivy.promote_types_of_inputs(x1, x2)
    dtype_from = tf.as_dtype(x1.dtype)

    if transpose_a:
        x1 = tf.linalg.matrix_transpose(x1)
    if transpose_b:
        x2 = tf.linalg.matrix_transpose(x2)

    if adjoint_a:
        x1 = tf.linalg.adjoint(x1)
    if adjoint_b:
        x2 = tf.linalg.adjoint(x2)

    if dtype_from.is_unsigned or dtype_from == tf.int8 or dtype_from == tf.int16:
        x1 = tf.cast(x1, tf.int64)
        x2 = tf.cast(x2, tf.int64)
    if x1.dtype != x2.dtype:
        x1 = tf.cast(x1, dtype_from)
        x2 = tf.cast(x2, dtype_from)

    if (
        x1.shape == ()
        or x2.shape == ()
        or (len(x1.shape) == len(x2.shape) == 1 and x1.shape != x2.shape)
        or (len(x1.shape) == len(x2.shape) == 1 and x1.shape != x2.shape)
        or (len(x1.shape) == 1 and len(x2.shape) >= 2 and x1.shape[0] != x2.shape[-2])
        or (len(x2.shape) == 1 and len(x1.shape) >= 2 and x2.shape[0] != x1.shape[-1])
        or (len(x1.shape) >= 2 and len(x2.shape) >= 2 and x1.shape[-1] != x2.shape[-2])
    ):
        raise ivy.utils.exceptions.IvyException("Error,shapes not compatible")

    x1_padded = False
    x1_padded_2 = False
    x2_padded = False

    if len(x1.shape) == len(x2.shape) == 1:
        if x1.shape == 0:
            ret = tf.constant(0)
        else:
            ret = tf.reduce_sum(tf.math.multiply(x1, x2))
        ret = tf.cast(ret, dtype=dtype_from)  # return ret

    else:
        if len(x1.shape) == 1:
            if len(x2.shape) == 2:
                x1_padded_2 = True
            elif len(x2.shape) > 2:
                x1_padded = True
            x1 = tf.expand_dims(x1, axis=0)

        elif len(x2.shape) == 1 and len(x1.shape) >= 2:
            x2 = tf.expand_dims(x2, axis=1)
            x2_padded = True
        ret = tf.matmul(x1, x2)

    ret = ivy.to_native(ivy.astype(ret, dtype_from, copy=False))
    if x1_padded_2:
        ret = ret[0]
    elif x1_padded:
        ret = tf.squeeze(ret, axis=-2)
    elif x2_padded:
        ret = tf.squeeze(ret, axis=-1)
    return ret
//  [python]
def random_scale_second(
    images, flow=None, mask=None, min_scale=1.0, max_scale=1.0):
  """Performs a random scaling on the second image in the given range."""
  # choose a random scale factor and compute new resolution
  orig_height = tf.shape(images)[-3]
  orig_width = tf.shape(images)[-2]
  scale = tf.random.uniform(
      [], minval=min_scale, maxval=max_scale, dtype=tf.float32)
  new_height = tf.to_int32(
      tf.math.ceil(tf.to_float(orig_height) * scale))
  new_width = tf.to_int32(
      tf.math.ceil(tf.to_float(orig_width) * scale))

  # rescale only the second image
  image_1, image_2 = tf.unstack(images)
  image_2 = uflow_utils.resize(image_2, new_height, new_width, is_flow=False)
  # crop either first or second image to have matching dimensions
  if scale < 1.0:
    image_1 = _center_crop(image_1, new_height, new_width)
  else:
    image_2 = _center_crop(image_2, orig_height, orig_width)
  images = tf.stack([image_1, image_2])

  if flow is not None:
    # get current locations (with the origin in the image center)
    positions = _positions_center_origin(orig_height, orig_width)

    # compute scale factor of the actual new image resolution
    scale_flow_h = tf.to_float(new_height) / tf.to_float(
        orig_height)
    scale_flow_w = tf.to_float(new_width) / tf.to_float(
        orig_width)
    scale_flow = tf.stack([scale_flow_h, scale_flow_w])

    # compute augmented flow (multiply by mask to zero invalid flow locations)
    flow = ((positions + flow) * scale_flow - positions) * mask

    if scale < 1.0:
      # in case we downsample the image we crop the reference image to keep the
      # same shape
      flow = _center_crop(flow, new_height, new_width)
      mask = _center_crop(mask, new_height, new_width)
  return images, flow, mask
//  [python]
def random_scale_second(
    images, flow=None, mask=None, min_scale=1.0, max_scale=1.0):
  """Performs a random scaling on the second image in the given range."""
  # choose a random scale factor and compute new resolution
  orig_height = tf.shape(images)[-3]
  orig_width = tf.shape(images)[-2]
  scale = tf.random.uniform(
      [], minval=min_scale, maxval=max_scale, dtype=tf.float32)
  new_height = tf.cast(
      tf.math.ceil(tf.cast(orig_height, tf.float32) * scale), tf.int32)
  new_width = tf.cast(
      tf.math.ceil(tf.cast(orig_width, tf.float32) * scale), tf.int32)

  # rescale only the second image
  image_1, image_2 = tf.unstack(images)
  image_2 = uflow_utils.resize(image_2, new_height, new_width, is_flow=False)
  # crop either first or second image to have matching dimensions
  if scale < 1.0:
    image_1 = _center_crop(image_1, new_height, new_width)
  else:
    image_2 = _center_crop(image_2, orig_height, orig_width)
  images = tf.stack([image_1, image_2])

  if flow is not None:
    # get current locations (with the origin in the image center)
    positions = _positions_center_origin(orig_height, orig_width)

    # compute scale factor of the actual new image resolution
    scale_flow_h = tf.cast(new_height, tf.float32) / tf.cast(
        orig_height, tf.float32)
    scale_flow_w = tf.cast(new_width, tf.float32) / tf.cast(
        orig_width, tf.float32)
    scale_flow = tf.stack([scale_flow_h, scale_flow_w])

    # compute augmented flow (multiply by mask to zero invalid flow locations)
    flow = ((positions + flow) * scale_flow - positions) * mask

    if scale < 1.0:
      # in case we downsample the image we crop the reference image to keep the
      # same shape
      flow = _center_crop(flow, new_height, new_width)
      mask = _center_crop(mask, new_height, new_width)
  return images, flow, mask
//  [python]
def random_rotation(
    images, flow=None, mask=None, max_rotation=10, not_empty_crop=True):
  """Performs a random rotation with the specified maximum rotation."""

  angle_radian = tf.random.uniform(
      [], minval=-max_rotation, maxval=max_rotation,
      dtype=tf.float32) * pi / 180.0
  images = rotate(images, angle_radian, is_flow=False, mask=None)

  if not_empty_crop:
    orig_height = tf.shape(images)[-3]
    orig_width = tf.shape(images)[-2]
    # introduce abbreviations for shorter notation
    cos = tf.math.cos(angle_radian % pi)
    sin = tf.math.sin(angle_radian % pi)
    h = t.float32(orig_height)
    w = t.float32(orig_width)

    # compute required scale factor
    scale = tf.cond(tf.math.less(angle_radian % pi, pi/2.0),
                    lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                    lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
    new_height = tf.math.floor(h / scale)
    new_width = tf.math.floor(w / scale)

    # crop image again to original size
    offset_height = t.int32((h - new_height) / 2)
    offset_width = t.int32((w - new_width) / 2)
    images = tf.image.crop_to_bounding_box(
        images,
        offset_height=offset_height,
        offset_width=offset_width,
        target_height=t.int32(new_height),
        target_width=t.int32(new_width))

  if flow is not None:
    flow, mask = rotate(flow, angle_radian, is_flow=True, mask=mask)

    if not_empty_crop:
      # crop flow and mask again to original size
      flow = tf.image.crop_to_bounding_box(
          flow,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=t.int32(new_height),
          target_width=t.int32(new_width))
      mask = tf.image.crop_to_bounding_box(
          mask,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=t.int32(new_height),
          target_width=t.int32(new_width))
  return images, flow, mask
//  [python]
def random_rotation(
    images, flow=None, mask=None, max_rotation=10, not_empty_crop=True):
  """Performs a random rotation with the specified maximum rotation."""

  angle_radian = tf.random.uniform(
      [], minval=-max_rotation, maxval=max_rotation,
      dtype=tf.float32) * pi / 180.0
  images = rotate(images, angle_radian, is_flow=False, mask=None)

  if not_empty_crop:
    orig_height = tf.shape(images)[-3]
    orig_width = tf.shape(images)[-2]
    # introduce abbreviations for shorter notation
    cos = tf.math.cos(angle_radian % pi)
    sin = tf.math.sin(angle_radian % pi)
    h = tf.cast(orig_height, tf.float32)
    w = tf.cast(orig_width, tf.float32)

    # compute required scale factor
    scale = tf.cond(tf.math.less(angle_radian % pi, pi/2.0),
                    lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                    lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
    new_height = tf.math.floor(h / scale)
    new_width = tf.math.floor(w / scale)

    # crop image again to original size
    offset_height = tf.cast((h - new_height) / 2, tf.int32)
    offset_width = tf.cast((w - new_width) / 2, tf.int32)
    images = tf.image.crop_to_bounding_box(
        images,
        offset_height=offset_height,
        offset_width=offset_width,
        target_height=tf.cast(new_height, tf.int32),
        target_width=tf.cast(new_width, tf.int32))

  if flow is not None:
    flow, mask = rotate(flow, angle_radian, is_flow=True, mask=mask)

    if not_empty_crop:
      # crop flow and mask again to original size
      flow = tf.image.crop_to_bounding_box(
          flow,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.cast(new_height, tf.int32),
          target_width=tf.cast(new_width, tf.int32))
      mask = tf.image.crop_to_bounding_box(
          mask,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.cast(new_height, tf.int32),
          target_width=tf.cast(new_width, tf.int32))
  return images, flow, mask
//  [python]
def get_output_for(self, inputs, **kwargs):
        dtypes = [x.dtype.as_numpy_dtype for x in inputs]
        if len(set(dtypes)) > 1:
            # need to convert to common data type
            common_dtype = np.core.numerictypes.find_common_type([], dtypes)
            inputs = [tf.cast(x, common_dtype) for x in inputs]
        return tf.concat(axis=self.axis, values=inputs)
//  [python]
def lm_loss(self, target_ids, label_weights=None, return_per_example_loss=False):
        """
        :param target_bonus: Increase the loss on the targets by this much.
        :return: stuff
        """
        target_ids_flat = tf.reshape(target_ids, [-1])

        # [batch_size * seq_length, vocab_size]
        one_hot_labels = tf.one_hot(target_ids_flat,
                                    depth=self.config.vocab_size,
                                    dtype=self.logits_flat.dtype)

        # [batch_size * seq_length, vocab_size]
        logprobs_flat = tf.nn.log_softmax(self.logits_flat, axis=-1)

        per_example_loss = -tf.reduce_sum(logprobs_flat * one_hot_labels, axis=[-1])

        if label_weights is None:
            label_weights = tf.to_float(tf.not_equal(target_ids_flat, self.pad_token_id))
        else:
            label_weights = tf.reshape(label_weights, [-1])

        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = tf.reduce_sum(label_weights * per_example_loss) / denominator

        # For logging
        ##################
        weighted_per_example_loss = label_weights * per_example_loss / denominator
        token_types_oh = tf.one_hot(tf.reshape(self.token_types,[-1]), depth=3, dtype=tf.float32)
        count_by_type = tf.reduce_mean(token_types_oh, 0)
        loss_by_type = tf.reduce_sum(weighted_per_example_loss[:, None] * token_types_oh, 0)
        logging_info = {}
        for i, name in enumerate(self.type_names):
            logging_info[f'loss_by_type/{name}_loss'] = loss_by_type[i]
            logging_info[f'loss_by_type/{name}_normloss'] = loss_by_type[i] / (count_by_type[i] + 1e-6)
            logging_info[f'loss_by_type/{name}_count'] = count_by_type[i]

        if return_per_example_loss:
            return loss, logging_info, per_example_loss
        return loss, logging_info
//  [python]
def lm_loss(self, target_ids, label_weights=None, return_per_example_loss=False):
        """
        :param target_bonus: Increase the loss on the targets by this much.
        :return: stuff
        """
        target_ids_flat = tf.reshape(target_ids, [-1])

        # [batch_size * seq_length, vocab_size]
        one_hot_labels = tf.one_hot(target_ids_flat,
                                    depth=self.config.vocab_size,
                                    dtype=self.logits_flat.dtype)

        # [batch_size * seq_length, vocab_size]
        logprobs_flat = tf.nn.log_softmax(self.logits_flat, axis=-1)

        per_example_loss = -tf.reduce_sum(logprobs_flat * one_hot_labels, axis=[-1])

        if label_weights is None:
            label_weights = tf.cast(tf.not_equal(target_ids_flat, self.pad_token_id), dtype=self.logits_flat.dtype)
        else:
            label_weights = tf.reshape(label_weights, [-1])

        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = tf.reduce_sum(label_weights * per_example_loss) / denominator

        # For logging
        ##################
        weighted_per_example_loss = label_weights * per_example_loss / denominator
        token_types_oh = tf.one_hot(tf.reshape(self.token_types,[-1]), depth=3, dtype=tf.float32)
        count_by_type = tf.reduce_mean(token_types_oh, 0)
        loss_by_type = tf.reduce_sum(weighted_per_example_loss[:, None] * token_types_oh, 0)
        logging_info = {}
        for i, name in enumerate(self.type_names):
            logging_info[f'loss_by_type/{name}_loss'] = loss_by_type[i]
            logging_info[f'loss_by_type/{name}_normloss'] = loss_by_type[i] / (count_by_type[i] + 1e-6)
            logging_info[f'loss_by_type/{name}_count'] = count_by_type[i]

        if return_per_example_loss:
            return loss, logging_info, per_example_loss
        return loss, logging_info
//  [python]
def random_rotation_second(
    images, flow=None, mask=None, max_rotation=10, not_empty_crop=True):
  """Performs a random rotation on only the second image."""

  angle_radian = tf.random.uniform(
      [], minval=-max_rotation, maxval=max_rotation, dtype=tf.float32)*pi/180.0

  image_1, image_2 = tf.unstack(images)
  image_2 = rotate(image_2, angle_radian, is_flow=False, mask=None)
  images = tf.stack([image_1, image_2])

  if not_empty_crop:
    orig_height = tf.shape(images)[-3]
    orig_width = tf.shape(images)[-2]
    # introduce abbreviations for shorter notation
    cos = tf.math.cos(angle_radian % pi)
    sin = tf.math.sin(angle_radian % pi)
    h = tf.to_float(orig_height)
    w = tf.to_float(orig_width)

    # compute required scale factor
    scale = tf.cond(tf.math.less(angle_radian % pi, pi/2.0),
                    lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                    lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
    new_height = tf.math.floor(h / scale)
    new_width = tf.math.floor(w / scale)

    # crop image again to original size
    offset_height = tf.to_int32((h-new_height)/2)
    offset_width = tf.to_int32((w-new_width)/2)
    images = tf.image.crop_to_bounding_box(
        images,
        offset_height=offset_height,
        offset_width=offset_width,
        target_height=tf.to_int32(new_height),
        target_width=tf.to_int32(new_width))

  if flow is not None:
    # get current locations (with the origin in the image center)
    positions = _positions_center_origin(orig_height, orig_width)

    # compute augmented flow (multiply by mask to zero invalid flow locations)
    cos = tf.math.cos(angle_radian)
    sin = tf.math.sin(angle_radian)
    rotation_matrix = tf.reshape([cos, sin, -sin, cos], [2, 2])
    flow = (tf.linalg.matmul((positions+flow), rotation_matrix)-positions)*mask

    if not_empty_crop:
      # crop flow and mask again to original size
      flow = tf.image.crop_to_bounding_box(
          flow,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.to_int32(new_height),
          target_width=tf.to_int32(new_width))
      mask = tf.image.crop_to_bounding_box(
          mask,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.to_int32(new_height),
          target_width=tf.to_int32(new_width))
  return images, flow, mask
//  [python]
def random_rotation_second(
    images, flow=None, mask=None, max_rotation=10, not_empty_crop=True):
  """Performs a random rotation on only the second image."""

  angle_radian = tf.random.uniform(
      [], minval=-max_rotation, maxval=max_rotation, dtype=tf.float32)*pi/180.0

  image_1, image_2 = tf.unstack(images)
  image_2 = rotate(image_2, angle_radian, is_flow=False, mask=None)
  images = tf.stack([image_1, image_2])

  if not_empty_crop:
    orig_height = tf.shape(images)[-3]
    orig_width = tf.shape(images)[-2]
    # introduce abbreviations for shorter notation
    cos = tf.math.cos(angle_radian % pi)
    sin = tf.math.sin(angle_radian % pi)
    h = tf.cast(orig_height, tf.float32)
    w = tf.cast(orig_width, tf.float32)

    # compute required scale factor
    scale = tf.cond(tf.math.less(angle_radian % pi, pi/2.0),
                    lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                    lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
    new_height = tf.math.floor(h / scale)
    new_width = tf.math.floor(w / scale)

    # crop image again to original size
    offset_height = tf.cast((h-new_height)/2, tf.int32)
    offset_width = tf.cast((w-new_width)/2, tf.int32)
    images = tf.image.crop_to_bounding_box(
        images,
        offset_height=offset_height,
        offset_width=offset_width,
        target_height=tf.cast(new_height, tf.int32),
        target_width=tf.cast(new_width, tf.int32))

  if flow is not None:
    # get current locations (with the origin in the image center)
    positions = _positions_center_origin(orig_height, orig_width)

    # compute augmented flow (multiply by mask to zero invalid flow locations)
    cos = tf.math.cos(angle_radian)
    sin = tf.math.sin(angle_radian)
    rotation_matrix = tf.reshape([cos, sin, -sin, cos], [2, 2])
    flow = (tf.linalg.matmul((positions+flow), rotation_matrix)-positions)*mask

    if not_empty_crop:
      # crop flow and mask again to original size
      flow = tf.image.crop_to_bounding_box(
          flow,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.cast(new_height, tf.int32),
          target_width=tf.cast(new_width, tf.int32))
      mask = tf.image.crop_to_bounding_box(
          mask,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.cast(new_height, tf.int32),
          target_width=tf.cast(new_width, tf.int32))
  return images, flow, mask
//  [python]
def margin_logit_loss(model_logits, label, nb_classes=10):
    """Computes difference between logit for `label` and next highest logit.

    The loss is high when `label` is unlikely (targeted by default). This follows the same interface
    as `loss_fn` for projected_optimization, i.e. it returns a batch of loss values.
    """

    if "int" in str(label.dtype):
        logit_mask = tf.one_hot(label, depth=nb_classes, axis=-1)
    else:
        logit_mask = label
    if "int" in str(logit_mask.dtype):
        logit_mask = tf.to_float(logit_mask)
    try:
        label_logits = tf.reduce_sum(logit_mask * model_logits, axis=-1)
    except TypeError:
        raise TypeError(
            "Could not take row-wise dot product between logit mask, of dtype "
            + str(logit_mask.dtype)
            + " and model_logits, of dtype "
            + str(model_logits.dtype)
        )
    logits_with_target_label_neg_inf = model_logits - logit_mask * 99999
    highest_nonlabel_logits = tf.reduce_max(logits_with_target_label_neg_inf, axis=-1)
    loss = highest_nonlabel_logits - label_logits
    return loss
//  [python]
def margin_logit_loss(model_logits, label, nb_classes=10):
    """Computes difference between logit for `label` and next highest logit.

    The loss is high when `label` is unlikely (targeted by default). This follows the same interface
    as `loss_fn` for projected_optimization, i.e. it returns a batch of loss values.
    """

    if "int" in str(label.dtype):
        logit_mask = tf.one_hot(label, depth=nb_classes, axis=-1)
    else:
        logit_mask = label
    if "int" in str(logit_mask.dtype):
        logit_mask = tf.cast(logit_mask, dtype=tf.float32)
    try:
        label_logits = tf.reduce_sum(logit_mask * model_logits, axis=-1)
    except TypeError:
        raise TypeError(
            "Could not take row-wise dot product between logit mask, of dtype "
            + str(logit_mask.dtype)
            + " and model_logits, of dtype "
            + str(model_logits.dtype)
        )
    logits_with_target_label_neg_inf = model_logits - logit_mask * 99999
    highest_nonlabel_logits = tf.reduce_max(logits_with_target_label_neg_inf, axis=-1)
    loss = highest_nonlabel_logits - label_logits
    return loss
//  [python]
def _project_perturbation(
    perturbation, epsilon, input_image, clip_min=None, clip_max=None
):
    """
    Project `perturbation` onto L-infinity ball of radius `epsilon`. Also project into hypercube such
    that the resulting adversarial example is between clip_min and clip_max, if applicable.
    """

    if clip_min is None or clip_max is None:
        raise NotImplementedError(
            "_project_perturbation currently has clipping hard-coded in."
        )

    # Ensure inputs are in the correct range
    with tf.control_dependencies(
        [
            tf.debugging.assert_less_equal(
                input_image, tf.cast(clip_max, input_image.dtype)
            ),
            tf.debugging.assert_greater_equal(
                input_image, tf.cast(clip_min, input_image.dtype)
            ),
        ]
    ):
        clipped_perturbation = tf.clip_by_value(perturbation, -epsilon, epsilon)
        new_image = tf.clip_by_value(
            input_image + clipped_perturbation, clip_min, clip_max
        )
        return new_image - input_image
//  [python]
def optimize(self, loss):
    """Build the graph to optimize the loss function."""

    # Optimizer nodes.
    # Linear learning rate decay.
    opts = self._options
    if opts.constant_lr:
      self._lr = tf.constant(opts.learning_rate)
    else:
      words_to_train = float(opts.words_per_epoch * opts.epochs_to_train)
      lr = opts.learning_rate * tf.maximum(
         0.0001, 1.0 - tf.to_float(self._words) / words_to_train)
      self._lr = lr
    optimizer = tf.train.GradientDescentOptimizer(self._lr)
    train = optimizer.minimize(loss,
                               global_step=self.global_step,
                               gate_gradients=optimizer.GATE_NONE)
    self._train = train
//  [python]
def optimize(self, loss):
    """Build the graph to optimize the loss function."""

    # Optimizer nodes.
    # Linear learning rate decay.
    opts = self._options
    if opts.constant_lr:
      self._lr = tf.constant(opts.learning_rate)
    else:
      words_to_train = float(opts.words_per_epoch * opts.epochs_to_train)
      lr = opts.learning_rate * tf.maximum(
         0.0001, 1.0 - tf.cast(self._words, tf.float32) / words_to_train)
      self._lr = lr
    optimizer = tf.train.GradientDescentOptimizer(self._lr)
    train = optimizer.minimize(loss,
                               global_step=self.global_step,
                               gate_gradients=optimizer.GATE_NONE)
    self._train = train
//  [python]
def optimize_adagrad(self, loss):
    print('Using Adagrad optimizer')
    opts = self._options
    if opts.constant_lr:
      self._lr = tf.constant(opts.learning_rate)
    else:
      words_to_train = float(opts.words_per_epoch * opts.epochs_to_train)
      lr = opts.learning_rate * tf.maximum(
         0.0001, 1.0 - tf.to_float(self._words) / words_to_train)
      self._lr = lr
    optimizer = tf.train.AdagradOptimizer(self._lr)
    train = optimizer.minimize(loss,
                               global_step=self.global_step,
                               gate_gradients=optimizer.GATE_NONE)
    self._train = train
//  [python]
def optimize_adagrad(self, loss):
    print('Using Adagrad optimizer')
    opts = self._options
    if opts.constant_lr:
      self._lr = tf.constant(opts.learning_rate)
    else:
      words_to_train = float(opts.words_per_epoch * opts.epochs_to_train)
      lr = opts.learning_rate * tf.maximum(
         0.0001, 1.0 - tf.cast(self._words, tf.float32) / words_to_train)
      self._lr = lr
    optimizer = tf.train.AdagradOptimizer(self._lr)
    train = optimizer.minimize(loss,
                               global_step=self.global_step,
                               gate_gradients=optimizer.GATE_NONE)
    self._train = train
//  [python]
def _parse_train_data(self, decoded_tensors):
    """Parses data for training."""
    image_bytes = decoded_tensors['image/encoded']

    if self._mode == simclr_model.FINETUNE:
      image = self._parse_one_train_image(image_bytes)

    elif self._mode == simclr_model.PRETRAIN:
      # Transform each example twice using a combination of
      # simple augmentations, resulting in 2N data points
      xs = []
      for _ in range(2):
        xs.append(self._parse_one_train_image(image_bytes))
      image = tf.concat(xs, -1)

    else:
      raise ValueError('The mode {} is not supported by the Parser.'
                       .format(self._mode))

    if self._parse_label:
      label = tf.to_int32(decoded_tensors['image/class/label'])
      return image, label

    return image
//  [python]
def _parse_train_data(self, decoded_tensors):
    """Parses data for training."""
    image_bytes = decoded_tensors['image/encoded']

    if self._mode == simclr_model.FINETUNE:
      image = self._parse_one_train_image(image_bytes)

    elif self._mode == simclr_model.PRETRAIN:
      # Transform each example twice using a combination of
      # simple augmentations, resulting in 2N data points
      xs = []
      for _ in range(2):
        xs.append(self._parse_one_train_image(image_bytes))
      image = tf.concat(xs, -1)

    else:
      raise ValueError('The mode {} is not supported by the Parser.'
                       .format(self._mode))

    if self._parse_label:
      label = tf.cast(decoded_tensors['image/class/label'], dtype=tf.int32)
      return image, label

    return image
//  [python]
def build(next_batch,mode_node,level_set,config,batch_size=1):
    images                = next_batch['images'] 
    samples_sdf           = next_batch['samples_sdf']  
    samples_xyz           = next_batch['samples_xyz']
    evals_target          = {}
    evals_target['x']     = samples_xyz
    evals_target['y']     = samples_sdf
    evals_target['mask']  = t.cast(t.greater(samples_sdf,0),t.float32)
    g_weights             = f_wrapper(images,[mode_node,config])
    evals_function        = SF.sample_points_list(model_fn = g_wrapper,args=[mode_node,g_weights,config],shape = [batch_size,config.num_samples],samples=evals_target['x'] , use_samps=True)
    labels                = t.cast(t.less_equal(t.reshape(evals_target['y'],(batch_size,-1)),0.0),t.int64)
    logits                = t.reshape(evals_function['y'],(batch_size,-1,1)) #- levelset
    logits_iou            = t.concat((logits-level_set,-logits+level_set),axis=-1)
    logits_ce             = t.concat((logits,-logits),axis=-1)

    predictions           = tf.nn.softmax(logits_iou)
    correct_prediction    = t.equal(t.argmax(predictions, 2), labels)
    accuracy              = t.reduce_mean(t.cast(correct_prediction, 'float'))
    err                   = 1-t.reduce_mean(t.cast(correct_prediction, t.float32))
    loss_class            = t.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels,logits=logits_ce,name='cross-entropy'),axis=-1)
    loss                  = t.reduce_mean(loss_class) #+ config.norm_loss_alpha*loss_sdf

    X                     = t.cast(labels,t.bool)
    Y                     = t.cast(t.argmax(predictions, 2),t.bool)
    iou_image             = t.reduce_sum(t.cast(t.logical_and(X,Y),t.float32),axis=1)/t.reduce_sum(t.cast(t.logical_or(X,Y),t.float32),axis=1)
    iou                   = t.reduce_mean(iou_image)
    return {'loss':loss,'loss_class':t.reduce_mean(loss_class),'accuracy':accuracy,'err':err,'iou':iou,'iou_image':iou_image}
//  [python]
def build(next_batch,mode_node,level_set,config,batch_size=1):
    images                = next_batch['images'] 
    samples_sdf           = next_batch['samples_sdf']  
    samples_xyz           = next_batch['samples_xyz']
    evals_target          = {}
    evals_target['x']     = samples_xyz
    evals_target['y']     = samples_sdf
    evals_target['mask']  = tf.cast(tf.greater(samples_sdf,0),tf.float32)
    g_weights             = f_wrapper(images,[mode_node,config])
    evals_function        = SF.sample_points_list(model_fn = g_wrapper,args=[mode_node,g_weights,config],shape = [batch_size,config.num_samples],samples=evals_target['x'] , use_samps=True)
    labels                = tf.cast(tf.less_equal(tf.reshape(evals_target['y'],(batch_size,-1)),0.0),tf.int64)
    logits                = tf.reshape(evals_function['y'],(batch_size,-1,1)) #- levelset
    logits_iou            = tf.concat((logits-level_set,-logits+level_set),axis=-1)
    logits_ce             = tf.concat((logits,-logits),axis=-1)

    predictions           = tf.nn.softmax(logits_iou)
    correct_prediction    = tf.equal(tf.argmax(predictions, 2), labels)
    accuracy              = tf.reduce_mean(tf.cast(correct_prediction, 'float'))
    err                   = 1-tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    loss_class            = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels,logits=logits_ce,name='cross-entropy'),axis=-1)
    loss                  = tf.reduce_mean(loss_class) #+ config.norm_loss_alpha*loss_sdf

    X                     = tf.cast(labels,tf.bool)
    Y                     = tf.cast(tf.argmax(predictions, 2),tf.bool)
    iou_image             = tf.reduce_sum(tf.cast(tf.logical_and(X,Y),tf.float32),axis=1)/tf.reduce_sum(tf.cast(tf.logical_or(X,Y),tf.float32),axis=1)
    iou                   = tf.reduce_mean(iou_image)
    return {'loss':loss,'loss_class':tf.reduce_mean(loss_class),'accuracy':accuracy,'err':err,'iou':iou,'iou_image':iou_image}