jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _tfval_to_tensor_jax_dtype(val: TfVal,
                               jax_dtype: DType | None = None,
                               memoize_constants=False) -> tuple[TfVal, DType]:
  """Converts a scalar, ndarray, or tf.Tensor to a tf.Tensor with proper type.

  If `jax_dtype` is missing, uses JAX typing rules.
  See README.md for details regarding 64-bit values.

  Args:
    val: a scalar, ndarray, tf.Tensor, or tf.Variable
    jax_dtype: an optional dtype to use. If missing, uses JAX type inference
      rules for constants.
    memoize_constants: whether to memoize TF constants. We can't do this
      everywhere, we may be outside of a conversion scope.

  Returns:
    a tuple with a tf.Tensor with the type as needed by JAX, and the JAX type.
  """
  if isinstance(val, (tf.Tensor, tf.Variable)):
    jax_dtype = jax_dtype or _to_jax_dtype(val.dtype)  # Give JAX a chance to pick the type
    conversion_dtype = _to_tf_dtype(jax_dtype)
    if conversion_dtype != val.dtype:  # May need to cast for 64-bit values
      return t.cast(val, conversion_dtype), jax_dtype
    else:
      return val, jax_dtype
  else:  # A constant
    jax_dtype = jax_dtype or xla.abstractify(val).dtype
    # TODO(document): We assume that the value of a constant does not
    # change through the scope of the function. But it may be an ndarray, ...
    # JAX has the same problem when generating HLO.
    const_key = (id(val), jax_dtype)
    # Since we use id(val) as a cache key, we have to make sure that we keep
    # the previous `val` alive. Otherwise, for a ndarray, it can get garbage
    # collected and reused for a different value, which would create correctness
    # issues. We keep the `val` alive by storing in the cache the pair
    # `(val, tf_val)`.
    # Only memoize non-scalars. JAX will lift all non-scalar constants as
    # Jaxpr consts, to the top level of the Jaxpr. This ensures that we see them
    # early, when entering the Jaxpr, so we create the tf.const early and its
    # scope is the entire Jaxpr.
    do_memoize = (memoize_constants and np.size(val) > 1 and _thread_local_state.constant_cache is not None)
    if do_memoize:
      _, tf_val = _thread_local_state.constant_cache.get(const_key, (None, None))
    else:
      tf_val = None
    if tf_val is None:
      conversion_dtype = _to_tf_dtype(jax_dtype)
      # The float0 type is not known to TF.
      if jax_dtype == dtypes.float0:
        val = np.zeros(np.shape(val), conversion_dtype.as_numpy_dtype)
      if hasattr(val, 'dtype') and dtypes.issubdtype(val.dtype, dtypes.extended):
        val = val.dtype._rules.physical_const(val)
      tf_val = tf.convert_to_tensor(val, dtype=conversion_dtype)
      if do_memoize:
        _thread_local_state.constant_cache[const_key] = (val, tf_val)
    return tf_val, jax_dtype
//  [python]
def _tfval_to_tensor_jax_dtype(val: TfVal,
                               jax_dtype: DType | None = None,
                               memoize_constants=False) -> tuple[TfVal, DType]:
  """Converts a scalar, ndarray, or tf.Tensor to a tf.Tensor with proper type.

  If `jax_dtype` is missing, uses JAX typing rules.
  See README.md for details regarding 64-bit values.

  Args:
    val: a scalar, ndarray, tf.Tensor, or tf.Variable
    jax_dtype: an optional dtype to use. If missing, uses JAX type inference
      rules for constants.
    memoize_constants: whether to memoize TF constants. We can't do this
      everywhere, we may be outside of a conversion scope.

  Returns:
    a tuple with a tf.Tensor with the type as needed by JAX, and the JAX type.
  """
  if isinstance(val, (tf.Tensor, tf.Variable)):
    jax_dtype = jax_dtype or _to_jax_dtype(val.dtype)  # Give JAX a chance to pick the type
    conversion_dtype = _to_tf_dtype(jax_dtype)
    if conversion_dtype != val.dtype:  # May need to cast for 64-bit values
      return tf.cast(val, conversion_dtype), jax_dtype
    else:
      return val, jax_dtype
  else:  # A constant
    jax_dtype = jax_dtype or xla.abstractify(val).dtype
    # TODO(document): We assume that the value of a constant does not
    # change through the scope of the function. But it may be an ndarray, ...
    # JAX has the same problem when generating HLO.
    const_key = (id(val), jax_dtype)
    # Since we use id(val) as a cache key, we have to make sure that we keep
    # the previous `val` alive. Otherwise, for a ndarray, it can get garbage
    # collected and reused for a different value, which would create correctness
    # issues. We keep the `val` alive by storing in the cache the pair
    # `(val, tf_val)`.
    # Only memoize non-scalars. JAX will lift all non-scalar constants as
    # Jaxpr consts, to the top level of the Jaxpr. This ensures that we see them
    # early, when entering the Jaxpr, so we create the tf.const early and its
    # scope is the entire Jaxpr.
    do_memoize = (memoize_constants and np.size(val) > 1 and _thread_local_state.constant_cache is not None)
    if do_memoize:
      _, tf_val = _thread_local_state.constant_cache.get(const_key, (None, None))
    else:
      tf_val = None
    if tf_val is None:
      conversion_dtype = _to_tf_dtype(jax_dtype)
      # The float0 type is not known to TF.
      if jax_dtype == dtypes.float0:
        val = np.zeros(np.shape(val), conversion_dtype.as_numpy_dtype)
      if hasattr(val, 'dtype') and dtypes.issubdtype(val.dtype, dtypes.extended):
        val = val.dtype._rules.physical_const(val)
      tf_val = tf.convert_to_tensor(val, dtype=conversion_dtype)
      if do_memoize:
        _thread_local_state.constant_cache[const_key] = (val, tf_val)
    return tf_val, jax_dtype
//  [python]
def _cond(index: TfVal, *operands: TfVal, branches: Sequence[core.ClosedJaxpr],
          linear: Sequence[bool]) -> Sequence[TfVal]:
  del linear
  # tf.cond needs lambdas with no arguments.
  branches_tf = [
      partial(_interpret_jaxpr, jaxpr, *operands,
              # Same name stack as the XLA translation of cond_p
              extra_name_stack=f"branch_{i}_fun")
      for i, jaxpr in enumerate(branches)
  ]
  # Same name stack as XLA translation of cond_p
  # Note: extend_name_stack is a contextmanager, which is callable as a decorator.
  branches_tf = list(map(source_info_util.extend_name_stack("cond"),  # type: ignore[arg-type]
      branches_tf))
  if len(branches) == 2:
    # `index` comes with tf.int32 type of casted boolean parameter.
    return tf.cond(tf.cast(index, tf.bool), branches_tf[1], branches_tf[0])
  else:
    return tf.switch_case(index, branches_tf)
//  [python]
def _top_k(operand: TfVal, k: int) -> tuple[TfVal, TfVal]:
  # Some types originally incompatible with tf.math.top_k can be promoted
  # to a compatible type without loss of precision.
  def promote_tf_dtype(tf_dtype):
    if tf_dtype in [tf.bool, tf.uint8, tf.uint16]:
      return tf.uint32
    if tf_dtype in [tf.int8, tf.int16]:
      return tf.int32
    if tf_dtype is tf.float16:
      return tf.float32
    return None

  conversion_dtype = promote_tf_dtype(operand.dtype)
  if not core.is_constant_dim(k):
    k_tf = _eval_shape((k,))[0]
    k_tf = t.cast(k_tf, tf.int32)  # TopK works only for int32
  else:
    k_tf = k
  if conversion_dtype:
    values, indices = tf.math.top_k(
        t.cast(operand, conversion_dtype), k=k_tf, sorted=True)
    return t.cast(values, operand.dtype), indices
  else:
    return tf.math.top_k(operand, k=k_tf, sorted=True)
//  [python]
def _top_k(operand: TfVal, k: int) -> tuple[TfVal, TfVal]:
  # Some types originally incompatible with tf.math.top_k can be promoted
  # to a compatible type without loss of precision.
  def promote_tf_dtype(tf_dtype):
    if tf_dtype in [tf.bool, tf.uint8, tf.uint16]:
      return tf.uint32
    if tf_dtype in [tf.int8, tf.int16]:
      return tf.int32
    if tf_dtype is tf.float16:
      return tf.float32
    return None

  conversion_dtype = promote_tf_dtype(operand.dtype)
  if not core.is_constant_dim(k):
    k_tf = _eval_shape((k,))[0]
    k_tf = tf.cast(k_tf, tf.int32)  # TopK works only for int32
  else:
    k_tf = k
  if conversion_dtype:
    values, indices = tf.math.top_k(
        tf.dtypes.cast(operand, conversion_dtype), k=k_tf, sorted=True)
    return tf.dtypes.cast(values, operand.dtype), indices
  else:
    return tf.math.top_k(operand, k=k_tf, sorted=True)
//  [python]
def _eigh(
    operand: TfVal,
    lower: bool,
    sort_eigenvalues: bool,
    subset_by_index: tuple,
    _in_avals,
    _out_aval,
):
  del sort_eigenvalues
  if operand.shape[-1] == 0:
    v, w = operand, tf.reshape(operand, _eval_shape(_in_avals[0].shape[:-1]))
  else:
    if not lower:
      operand = tf.linalg.adjoint(operand)
    w, v = tf.linalg.eigh(operand)
  cast_type = {
      tf.complex64: tf.float32,
      tf.complex128: tf.float64
  }.get(operand.dtype)
  if not (subset_by_index is None or subset_by_index == (0, operand.shape[-1])):
    raise NotImplementedError("subset_by_index is not implemented")
  if cast_type is not None:
    w = t.cast(w, cast_type)
  return v, w
//  [python]
def _eigh(
    operand: TfVal,
    lower: bool,
    sort_eigenvalues: bool,
    subset_by_index: tuple,
    _in_avals,
    _out_aval,
):
  del sort_eigenvalues
  if operand.shape[-1] == 0:
    v, w = operand, tf.reshape(operand, _eval_shape(_in_avals[0].shape[:-1]))
  else:
    if not lower:
      operand = tf.linalg.adjoint(operand)
    w, v = tf.linalg.eigh(operand)
  cast_type = {
      tf.complex64: tf.float32,
      tf.complex128: tf.float64
  }.get(operand.dtype)
  if not (subset_by_index is None or subset_by_index == (0, operand.shape[-1])):
    raise NotImplementedError("subset_by_index is not implemented")
  if cast_type is not None:
    w = tf.cast(w, cast_type)
  return v, w
//  [python]
def handle_boolean_args(f, argnums: Sequence[int], boolean_f=None):
  """Computes functions with some bool args and bool results using int8.

  This is needed because some TF ops do not work for bool args, e.g.,
  inequalities, min/max.

  Args:
    f: a TF callable to wrap. It will be called with non-boolean arguments.
    argnums: the positional arguments that may be booleans.
    boolean_f: [Optional] a TF callable compatible with boolean
      arguments.

  Returns: a TF callable that can take a mix of boolean positional arguments
    (in the positions specified by `argnums`) and some non-boolean positional
    arguments. If there are no boolean arguments, just calls `f`. Otherwise,
    it calls `boolean_f` if defined. Otherwise, casts the boolean
    arguments to `int8`, calls `f`, then casts the result to `bool`.
  """
  argnums = tf.nest.flatten(argnums)

  def wrapper(*args: TfVal, **kwargs):
    argnum_types = {args[i].dtype for i in argnums}
    if tf.bool not in argnum_types:
      return f(*args, **kwargs)
    else:
      # All argnums should be boolean
      assert len(argnum_types) == 1, argnum_types
      if boolean_f != None:
        return boolean_f(*args, **kwargs)
      else:
        args_cast = [(t.cast(a, t.int8) if i in argnums else a)
                    for i, a in enumerate(args)]
        if "_in_avals" in kwargs:

          def cast_aval(aval):
            assert aval.dtype == np.bool_
            return core.ShapedArray(aval.shape, np.int8)

          _in_avals_cast = [
              cast_aval(aval) if i in argnums else aval
              for i, aval in enumerate(kwargs["_in_avals"])
          ]
          _out_aval_cast = tf.nest.map_structure(cast_aval, kwargs["_out_aval"])
          kwargs = dict(
              kwargs, _in_avals=_in_avals_cast, _out_aval=_out_aval_cast)
        out = f(*args_cast, **kwargs)
        return tf.nest.map_structure(lambda o: t.cast(o, t.bool), out)

  return wrapper
//  [python]
def handle_boolean_args(f, argnums: Sequence[int], boolean_f=None):
  """Computes functions with some bool args and bool results using int8.

  This is needed because some TF ops do not work for bool args, e.g.,
  inequalities, min/max.

  Args:
    f: a TF callable to wrap. It will be called with non-boolean arguments.
    argnums: the positional arguments that may be booleans.
    boolean_f: [Optional] a TF callable compatible with boolean
      arguments.

  Returns: a TF callable that can take a mix of boolean positional arguments
    (in the positions specified by `argnums`) and some non-boolean positional
    arguments. If there are no boolean arguments, just calls `f`. Otherwise,
    it calls `boolean_f` if defined. Otherwise, casts the boolean
    arguments to `int8`, calls `f`, then casts the result to `bool`.
  """
  argnums = tf.nest.flatten(argnums)

  def wrapper(*args: TfVal, **kwargs):
    argnum_types = {args[i].dtype for i in argnums}
    if tf.bool not in argnum_types:
      return f(*args, **kwargs)
    else:
      # All argnums should be boolean
      assert len(argnum_types) == 1, argnum_types
      if boolean_f != None:
        return boolean_f(*args, **kwargs)
      else:
        args_cast = [(tf.cast(a, tf.int8) if i in argnums else a)
                    for i, a in enumerate(args)]
        if "_in_avals" in kwargs:

          def cast_aval(aval):
            assert aval.dtype == np.bool_
            return core.ShapedArray(aval.shape, np.int8)

          _in_avals_cast = [
              cast_aval(aval) if i in argnums else aval
              for i, aval in enumerate(kwargs["_in_avals"])
          ]
          _out_aval_cast = tf.nest.map_structure(cast_aval, kwargs["_out_aval"])
          kwargs = dict(
              kwargs, _in_avals=_in_avals_cast, _out_aval=_out_aval_cast)
        out = f(*args_cast, **kwargs)
        return tf.nest.map_structure(lambda o: tf.cast(o, tf.bool), out)

  return wrapper
//  [python]
def creat_padding_mask(inputs):
    """
    对input中的padding单位进行mask
    :param inputs: 句子序列输入
    :return: 填充部分标记
    """
    mask = tf.to_float(tf.math.equal(inputs, 0))
    return mask[:, tf.newaxis, tf.newaxis, :]
//  [python]
def creat_padding_mask(inputs):
    """
    对input中的padding单位进行mask
    :param inputs: 句子序列输入
    :return: 填充部分标记
    """
    mask = tf.cast(tf.math.equal(inputs, 0), dtype=tf.float32)
    return mask[:, tf.newaxis, tf.newaxis, :]
//  [python]
def positional_encoding(position, deep):
    i = tf.range(deep, dtype=tf.float32)[tf.newaxis, :]
    position = tf.range(position, dtype=tf.float32)[:, tf.newaxis]
    angles = 1 / tf.pow(10000, (2 * (i // 2)) / tf.cast(deep, tf.float32))
    angle_rads = position * angles

    sines = tf.math.sin(angle_rads[:, 0::2])
    cosines = tf.math.cos(angle_rads[:, 1::2])
    pos_encoding = tf.concat([sines, cosines], axis=-1)
    pos_encoding = pos_encoding[tf.newaxis, ...]
    return tf.to_float(pos_encoding)
//  [python]
def positional_encoding(position, deep):
    i = tf.range(deep, dtype=tf.float32)[tf.newaxis, :]
    position = tf.range(position, dtype=tf.float32)[:, tf.newaxis]
    angles = 1 / tf.pow(10000, (2 * (i // 2)) / tf.cast(deep, tf.float32))
    angle_rads = position * angles

    sines = tf.math.sin(angle_rads[:, 0::2])
    cosines = tf.math.cos(angle_rads[:, 1::2])
    pos_encoding = tf.concat([sines, cosines], axis=-1)
    pos_encoding = pos_encoding[tf.newaxis, ...]
    return tf.cast(pos_encoding, tf.float32)
//  [python]
def __init__(self, d_model, warmup_steps=2000):
        super(CustomSchedule, self).__init__()
        self.d_model = d_model
        self.d_model = tf.cast(self.d_model, tf.float32)
        self.warmup_steps = warmup_steps
//  [python]
def encoder(vocab_size, deep, dropout):
    inputs = tf.keras.Input(shape=(None,))
    embedding = tf.keras.layers.Embedding(vocab_size, deep)(inputs)
    embedding *= tf.math.sqrt(tf.cast(deep, tf.float32))
    embedding = positional_encoding_layer(vocab_size, deep)(embedding)
    outputs = tf.keras.layers.Dropout(rate=dropout)(embedding)

    return tf.keras.Model(inputs=inputs, outputs=outputs)
//  [python]
def self_attention(query, key, value, mask):
    matmul = tf.matmul(query, key, transpose_b=True)
    deep = tf.to_float(tf.shape(key)[-1])
    scaled_attention = matmul / tf.math.sqrt(deep)

    if mask is not None:
        scaled_attention += (mask * -1e9)
    attention_weight = tf.nn.softmax(scaled_attention, axis=-1)
    output = tf.matmul(attention_weight, value)
    return output
//  [python]
def self_attention(query, key, value, mask):
    matmul = tf.matmul(query, key, transpose_b=True)
    deep = tf.cast(tf.shape(key)[-1], tf.float32)
    scaled_attention = matmul / tf.math.sqrt(deep)

    if mask is not None:
        scaled_attention += (mask * -1e9)
    attention_weight = tf.nn.softmax(scaled_attention, axis=-1)
    output = tf.matmul(attention_weight, value)
    return output
//  [python]
def train(dataset, epochs):
    for epoch in range(epochs):
        start = time.time()
        i = 0
        D_loss_list, G_loss_list = [], []
        for image_batch,target in dataset:
            i += 1
            img = t.cast(image_batch, tf.float32)
            imgs = normalization(img)
            train_step(imgs,target)
        print(epoch)        
        display.clear_output(wait=True)
        generate_and_save_images(conditional_gen,
                              epoch + 1,
                              seed)

        conditional_gen.save_weights('rock-paper-scissors/training_weights/gen_'+ str(epoch)+'.h5')
        conditional_discriminator.save_weights('rock-paper-scissors/training_weights/disc_'+ str(epoch)+'.h5')    
        print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

    # Generate after the final epoch
    display.clear_output(wait=True)
    generate_and_save_images(conditional_gen,
                            epochs,
                            seed)
//  [python]
def train(dataset, epochs):
    for epoch in range(epochs):
        start = time.time()
        i = 0
        D_loss_list, G_loss_list = [], []
        for image_batch,target in dataset:
            i += 1
            img = tf.cast(image_batch, tf.float32)
            imgs = normalization(img)
            train_step(imgs,target)
        print(epoch)        
        display.clear_output(wait=True)
        generate_and_save_images(conditional_gen,
                              epoch + 1,
                              seed)

        conditional_gen.save_weights('rock-paper-scissors/training_weights/gen_'+ str(epoch)+'.h5')
        conditional_discriminator.save_weights('rock-paper-scissors/training_weights/disc_'+ str(epoch)+'.h5')    
        print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

    # Generate after the final epoch
    display.clear_output(wait=True)
    generate_and_save_images(conditional_gen,
                            epochs,
                            seed)
//  [python]
def call(self, pixel_values: tf.Tensor, bool_masked_pos: tf.Tensor | None = None) -> tf.Tensor:
        embeddings = self.patch_embeddings(pixel_values)
        batch_size, seq_len, projection_dim = shape_list(embeddings)

        cls_tokens = tf.tile(self.cls_token, (batch_size, 1, 1))

        if bool_masked_pos is not None:
            mask_tokens = tf.broadcast_to(self.mask_token, (batch_size, seq_len, projection_dim))
            # replace the masked visual tokens by mask_tokens
            w = bool_masked_pos[..., None]
            w = tf.to_float(w)
            # since TF doesn't support eager tensor assignment
            embeddings = embeddings * (1 - w) + mask_tokens * w

        embeddings = tf.concat([cls_tokens, embeddings], axis=1)
        if self.position_embeddings is not None:
            embeddings = embeddings + self.position_embeddings
        embeddings = self.dropout(embeddings)

        return embeddings
//  [python]
def call(self, pixel_values: tf.Tensor, bool_masked_pos: tf.Tensor | None = None) -> tf.Tensor:
        embeddings = self.patch_embeddings(pixel_values)
        batch_size, seq_len, projection_dim = shape_list(embeddings)

        cls_tokens = tf.tile(self.cls_token, (batch_size, 1, 1))

        if bool_masked_pos is not None:
            mask_tokens = tf.broadcast_to(self.mask_token, (batch_size, seq_len, projection_dim))
            # replace the masked visual tokens by mask_tokens
            w = bool_masked_pos[..., None]
            w = tf.cast(w, mask_tokens.dtype)
            # since TF doesn't support eager tensor assignment
            embeddings = embeddings * (1 - w) + mask_tokens * w

        embeddings = tf.concat([cls_tokens, embeddings], axis=1)
        if self.position_embeddings is not None:
            embeddings = embeddings + self.position_embeddings
        embeddings = self.dropout(embeddings)

        return embeddings
//  [python]
def get_reset_goal(self, sampled_experience, cur_state, next_goal):
    # value function
    self._reset_goal_values = self.get_value_function(sampled_experience,
                                                      cur_state, next_goal)
    self._reset_goal_candidates = sampled_experience.observation[:, :self
                                                                 ._goal_dim]

    # distance to the initial state distribution
    self._distance_to_initial_state = t.cast(
        self._distance_fn.distance(self._reset_goal_candidates),
        dtype=t.float32)

    reset_goal_idx = tf.argmin(self._distance_to_initial_state -
                               self._lagrange_multiplier *
                               self._reset_goal_values)
    reset_goal = self._reset_goal_candidates[reset_goal_idx]
    self._values.append(self._reset_goal_values[reset_goal_idx])
    return reset_goal
//  [python]
def get_reset_goal(self, sampled_experience, cur_state, next_goal):
    # value function
    self._reset_goal_values = self.get_value_function(sampled_experience,
                                                      cur_state, next_goal)
    self._reset_goal_candidates = sampled_experience.observation[:, :self
                                                                 ._goal_dim]

    # distance to the initial state distribution
    self._distance_to_initial_state = tf.cast(
        self._distance_fn.distance(self._reset_goal_candidates),
        dtype=tf.float32)

    reset_goal_idx = tf.argmin(self._distance_to_initial_state -
                               self._lagrange_multiplier *
                               self._reset_goal_values)
    reset_goal = self._reset_goal_candidates[reset_goal_idx]
    self._values.append(self._reset_goal_values[reset_goal_idx])
    return reset_goal
//  [python]
def _efficientnet_preprocess(self, image):
    assert image is not None, "image cannot be None"
    assert self.output_height == self.output_width, "image output_width must be equal with output_height"
    image_size = self.output_height
    shape = tf.shape(image)
    image_height = shape[0]
    image_width = shape[1]

    padded_center_crop_size = tf.to_int32(
      ((image_size / (image_size + CROP_PADDING)) *
       tf.to_float(tf.minimum(image_height, image_width))))

    offset_height = ((image_height - padded_center_crop_size) + 1) // 2
    offset_width = ((image_width - padded_center_crop_size) + 1) // 2
    crop_window = tf.stack([
      offset_height, offset_width, padded_center_crop_size,
      padded_center_crop_size
    ])

    image = image[offset_height:offset_height + padded_center_crop_size, \
                 offset_width:offset_width + padded_center_crop_size, :]
    image = tf.image.resize_bicubic([image], [image_size, image_size])[0]

    image = self._mean_image_subtraction(image, EFFICIENTNET_MEAN_RGB)
    image = self._stddev_image_division(image, EFFICIENTNET_STDDEV_RGB)
    image = tf.expand_dims(image, 0)
    return image
//  [python]
def _efficientnet_preprocess(self, image):
    assert image is not None, "image cannot be None"
    assert self.output_height == self.output_width, "image output_width must be equal with output_height"
    image_size = self.output_height
    shape = tf.shape(image)
    image_height = shape[0]
    image_width = shape[1]

    padded_center_crop_size = tf.cast(
      ((image_size / (image_size + CROP_PADDING)) *
       tf.cast(tf.minimum(image_height, image_width), tf.float32)), tf.int32)

    offset_height = ((image_height - padded_center_crop_size) + 1) // 2
    offset_width = ((image_width - padded_center_crop_size) + 1) // 2
    crop_window = tf.stack([
      offset_height, offset_width, padded_center_crop_size,
      padded_center_crop_size
    ])

    image = image[offset_height:offset_height + padded_center_crop_size, \
                 offset_width:offset_width + padded_center_crop_size, :]
    image = tf.image.resize_bicubic([image], [image_size, image_size])[0]

    image = self._mean_image_subtraction(image, EFFICIENTNET_MEAN_RGB)
    image = self._stddev_image_division(image, EFFICIENTNET_STDDEV_RGB)
    image = tf.expand_dims(image, 0)
    return image
//  [python]
def compute_loss(self, logits, auxiliary_logits, labels):
        # upsample logits to the images' original size
        if len(shape_list(labels)) > 3:
            label_interp_shape = shape_list(labels)[1:-1]
        else:
            label_interp_shape = shape_list(labels)[-2:]

        upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
        if auxiliary_logits is not None:
            upsampled_auxiliary_logits = tf.image.resize(auxiliary_logits, size=label_interp_shape, method="bilinear")
        # compute weighted loss
        loss_fct = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")

        # Copied from https://www.tensorflow.org/text/tutorials/transformer#loss_and_metrics.
        # Utility to mask the index to ignore during computing the loss.
        def masked_loss(real, pred):
            mask = tf.math.logical_not(tf.math.equal(real, self.config.semantic_loss_ignore_index))
            loss_ = loss_fct(real, pred)
            mask = tf.to_float(mask)
            loss_ *= mask
            reduced_masked_loss = tf.reduce_sum(loss_) / tf.reduce_sum(mask)
            return tf.reshape(reduced_masked_loss, (1,))

        main_loss = masked_loss(labels, upsampled_logits)
        auxiliary_loss = masked_loss(labels, upsampled_auxiliary_logits)
        loss = main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

        return loss
//  [python]
def compute_loss(self, logits, auxiliary_logits, labels):
        # upsample logits to the images' original size
        if len(shape_list(labels)) > 3:
            label_interp_shape = shape_list(labels)[1:-1]
        else:
            label_interp_shape = shape_list(labels)[-2:]

        upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
        if auxiliary_logits is not None:
            upsampled_auxiliary_logits = tf.image.resize(auxiliary_logits, size=label_interp_shape, method="bilinear")
        # compute weighted loss
        loss_fct = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")

        # Copied from https://www.tensorflow.org/text/tutorials/transformer#loss_and_metrics.
        # Utility to mask the index to ignore during computing the loss.
        def masked_loss(real, pred):
            mask = tf.math.logical_not(tf.math.equal(real, self.config.semantic_loss_ignore_index))
            loss_ = loss_fct(real, pred)
            mask = tf.cast(mask, dtype=loss_.dtype)
            loss_ *= mask
            reduced_masked_loss = tf.reduce_sum(loss_) / tf.reduce_sum(mask)
            return tf.reshape(reduced_masked_loss, (1,))

        main_loss = masked_loss(labels, upsampled_logits)
        auxiliary_loss = masked_loss(labels, upsampled_auxiliary_logits)
        loss = main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

        return loss
//  [python]
def _pad_infos_object(self, image):
    """Get a Tensor to pad the info object list."""
    shape_ = tf.shape(image)
    val = tf.stack([
        t.cast(shape_[:2], tf.float32),
        t.cast(shape_[:2], tf.float32),
        t.ones_like(t.cast(shape_[:2], tf.float32)),
        t.zeros_like(t.cast(shape_[:2], tf.float32)),
    ])
    return val
//  [python]
def _pad_infos_object(self, image):
    """Get a Tensor to pad the info object list."""
    shape_ = tf.shape(image)
    val = tf.stack([
        tf.cast(shape_[:2], tf.float32),
        tf.cast(shape_[:2], tf.float32),
        tf.ones_like(tf.cast(shape_[:2], tf.float32)),
        tf.zeros_like(tf.cast(shape_[:2], tf.float32)),
    ])
    return val
//  [python]
def _parse_eval_data(self, data):
    """Parses data for evaluation."""

    # Get the image shape constants and cast the image to the selcted datatype.
    image = tf.cast(data['image'], dtype=self._dtype)
    boxes = data['groundtruth_boxes']
    classes = data['groundtruth_classes']

    image, infos, _ = preprocessing_ops.resize_and_jitter_image(
        image, [self._image_h, self._image_w],
        letter_box=self._letter_box,
        random_pad=False,
        shiftx=0.5,
        shifty=0.5,
        jitter=0.0)

    # Clip and clean boxes.
    image = image / 255.0
    boxes, inds = preprocessing_ops.transform_and_clip_boxes(
        boxes, infos, shuffle_boxes=False, area_thresh=0.0,
        filter_and_clip_boxes=False)
    classes = tf.gather(classes, inds)
    info = infos[-1]

    image, labels = self._build_label(
        image, boxes, classes, info, inds, data, is_training=False)
    return image, labels
//  [python]
def _create_model(self):
        # input features
        with tf.variable_scope('input_layer',
                               partitioner=self._input_layer_partitioner,
                               reuse=tf.AUTO_REUSE):
            fm_cols = {}
            if self._adaptive_emb and not self.tf:
                '''Adaptive Embedding Feature Part 1 of 2'''
                adaptive_mask_tensors = {}
                for col in CATEGORICAL_COLUMNS:
                    adaptive_mask_tensors[col] = tf.ones([args.batch_size],
                                                         tf.int32)
                dnn_input = tf.feature_column.input_layer(
                    features=self._feature,
                    feature_columns=self._deep_column,
                    adaptive_mask_tensors=adaptive_mask_tensors)
                wide_input = tf.feature_column.input_layer(
                self._feature, self._wide_column, cols_to_output_tensors=fm_cols, adaptive_mask_tensors=adaptive_mask_tensors)
            else:
                dnn_input = tf.feature_column.input_layer(self._feature,
                                                          self._deep_column)
                wide_input = tf.feature_column.input_layer(
                self._feature, self._wide_column, cols_to_output_tensors=fm_cols)

            fm_input = tf.stack([fm_cols[cols] for cols in self._fm_column], 1)

        if self.bf16:
            wide_input = tf.cast(wide_input, dtype=tf.bfloat16)
            fm_input = tf.cast(fm_input, dtype=tf.bfloat16)
            dnn_input = tf.cast(dnn_input, dtype=tf.bfloat16)

        # DNN part
        dnn_scope = tf.variable_scope('dnn')
        with dnn_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                else dnn_scope:
            dnn_output = self._dnn(dnn_input, self._dnn_hidden_units,
                                   'dnn_layer')

        # linear / fisrt order part
        with tf.variable_scope('linear', reuse=tf.AUTO_REUSE) as linear:
            linear_output = tf.reduce_sum(wide_input, axis=1, keepdims=True)

        # FM second order part
        with tf.variable_scope('fm', reuse=tf.AUTO_REUSE) as fm:
            sum_square = tf.square(tf.reduce_sum(fm_input, axis=1))
            square_sum = tf.reduce_sum(tf.square(fm_input), axis=1)
            fm_output = 0.5 * tf.subtract(sum_square, square_sum)

        # Final dnn layer
        all_input = tf.concat([dnn_output, linear_output, fm_output], 1)
        final_dnn_scope = tf.variable_scope('final_dnn')
        with final_dnn_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                else final_dnn_scope:
            dnn_logits = self._dnn(all_input, self._final_hidden_units, 'final_dnn')

        if self.bf16:
            dnn_logits = tf.cast(dnn_logits, dtype=tf.float32)

        self._logits = tf.layers.dense(dnn_logits, units=1)
        self.probability = tf.math.sigmoid(self._logits)
        self.output = tf.round(self.probability)
//  [python]
def get_padding_mask_from_valid_lengths(
    valid_lengths: tf.Tensor,
    max_length: Optional[int] = None,
    dtype: tf.dtypes.DType = tf.bool) -> tf.Tensor:
  """Gets a 2D mask of the padded region from valid lengths.

  Args:
    valid_lengths: A 1D int tensor containing the valid length of each row.
    max_length: (optional, int) The maximum length of each row. If `None`, the
      maximum value in `valid_lengths` will be used.
    dtype: The output dtype.

  Returns:
    2D padded region mask.
  """
  with tf.name_scope('get_padding_mask_from_valid_lengths'):
    if max_length is None:
      max_length = tf.reduce_max(valid_lengths)
    padding_mask = tf.logical_not(tf.sequence_mask(valid_lengths, max_length))

    return tf.to_float(padding_mask)
//  [python]
def get_padding_mask_from_valid_lengths(
    valid_lengths: tf.Tensor,
    max_length: Optional[int] = None,
    dtype: tf.dtypes.DType = tf.bool) -> tf.Tensor:
  """Gets a 2D mask of the padded region from valid lengths.

  Args:
    valid_lengths: A 1D int tensor containing the valid length of each row.
    max_length: (optional, int) The maximum length of each row. If `None`, the
      maximum value in `valid_lengths` will be used.
    dtype: The output dtype.

  Returns:
    2D padded region mask.
  """
  with tf.name_scope('get_padding_mask_from_valid_lengths'):
    if max_length is None:
      max_length = tf.reduce_max(valid_lengths)
    padding_mask = tf.logical_not(tf.sequence_mask(valid_lengths, max_length))

    return tf.cast(padding_mask, dtype=dtype)
//  [python]
def rotate_rboxes90(rboxes: tf.Tensor,
                    image_width: int,
                    image_height: int,
                    rotation_count: int = 1) -> tf.Tensor:
  """Rotate oriented rectangles counter-clockwise by multiples of 90 degrees."""
  image_width = t.cast(image_width, dtype=tf.float32)
  image_height = t.cast(image_height, dtype=tf.float32)

  rotation_count = rotation_count % 4
  x, y, w, h, angle = tf.split(rboxes, 5, axis=1)

  if rotation_count == 0:
    return rboxes
  elif rotation_count == 1:
    angle = tf.where(angle < -90.0, angle + 270, angle - 90)
    return tf.concat([y, image_width - x - 1, w, h, angle], axis=1)
  elif rotation_count == 2:
    angle = tf.where(angle < 0.0, angle + 180, angle - 180)
    return tf.concat([image_width - x - 1, image_height - y - 1, w, h, angle],
                     axis=1)
  else:
    angle = tf.where(angle > 90.0, angle - 270, angle + 90)
    return tf.concat([image_height - y - 1, x, w, h, angle], axis=1)
//  [python]
def rotate_rboxes90(rboxes: tf.Tensor,
                    image_width: int,
                    image_height: int,
                    rotation_count: int = 1) -> tf.Tensor:
  """Rotate oriented rectangles counter-clockwise by multiples of 90 degrees."""
  image_width = tf.cast(image_width, dtype=tf.float32)
  image_height = tf.cast(image_height, dtype=tf.float32)

  rotation_count = rotation_count % 4
  x, y, w, h, angle = tf.split(rboxes, 5, axis=1)

  if rotation_count == 0:
    return rboxes
  elif rotation_count == 1:
    angle = tf.where(angle < -90.0, angle + 270, angle - 90)
    return tf.concat([y, image_width - x - 1, w, h, angle], axis=1)
  elif rotation_count == 2:
    angle = tf.where(angle < 0.0, angle + 180, angle - 180)
    return tf.concat([image_width - x - 1, image_height - y - 1, w, h, angle],
                     axis=1)
  else:
    angle = tf.where(angle > 90.0, angle - 270, angle + 90)
    return tf.concat([image_height - y - 1, x, w, h, angle], axis=1)
//  [python]
def normalize_image_to_range(image: tf.Tensor,
                             original_minval: int = 0,
                             original_maxval: int = 255,
                             target_minval: float = -1.0,
                             target_maxval: float = 1.0) -> tf.Tensor:
  """Normalizes pixel values in the image.

  Moves the pixel values from the current [original_minval, original_maxval]
  range to the [target_minval, target_maxval] range.

  Args:
    image: A tensor of shape [height, width, channels]. Input will be converted
      to float32 type before normalization.
    original_minval: current image minimum value.
    original_maxval: current image maximum value.
    target_minval: target image minimum value.
    target_maxval: target image maximum value.

  Returns:
    A float tensor with the same shape as the input image.
  """
  if image.dtype is not tf.float32:
    image = t.cast(image, dtype=tf.float32)

  original_minval = float(original_minval)
  original_maxval = float(original_maxval)
  target_minval = float(target_minval)
  target_maxval = float(target_maxval)
  image = t.cast(image, dtype=tf.float32)
  image = tf.subtract(image, original_minval)
  image = tf.multiply(image, (target_maxval - target_minval) /
                      (original_maxval - original_minval))
  image = tf.add(image, target_minval)

  return image
//  [python]
def normalize_image_to_range(image: tf.Tensor,
                             original_minval: int = 0,
                             original_maxval: int = 255,
                             target_minval: float = -1.0,
                             target_maxval: float = 1.0) -> tf.Tensor:
  """Normalizes pixel values in the image.

  Moves the pixel values from the current [original_minval, original_maxval]
  range to the [target_minval, target_maxval] range.

  Args:
    image: A tensor of shape [height, width, channels]. Input will be converted
      to float32 type before normalization.
    original_minval: current image minimum value.
    original_maxval: current image maximum value.
    target_minval: target image minimum value.
    target_maxval: target image maximum value.

  Returns:
    A float tensor with the same shape as the input image.
  """
  if image.dtype is not tf.float32:
    image = tf.cast(image, dtype=tf.float32)

  original_minval = float(original_minval)
  original_maxval = float(original_maxval)
  target_minval = float(target_minval)
  target_maxval = float(target_maxval)
  image = tf.cast(image, dtype=tf.float32)
  image = tf.subtract(image, original_minval)
  image = tf.multiply(image, (target_maxval - target_minval) /
                      (original_maxval - original_minval))
  image = tf.add(image, target_minval)

  return image
//  [python]
def adjacency_loss(y_true, y_pred):
    """Implements Smooth-L1 loss.
    y_true and y_pred are typically: [N, 4], but could be any shape.
    """
    print(y_true.shape)
    print(y_pred.shape)
    y_true = tf.to_int32(y_true)
    total_count = tf.to_float(tf.size(y_true))
    count_ones = tf.to_float(tf.math.count_nonzero(y_true)) + tf.constant([1.0], tf.float32)
    count_zeros = total_count - count_ones
    pos_weight = tf.math.divide(count_zeros, count_ones)
    #loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=tf.one_hot(y_true, depth=2), logits=y_pred)
    loss = tf.nn.weighted_cross_entropy_with_logits(logits=y_pred,
                                                    targets=tf.one_hot(
                                                        y_true, depth=2),
                                                    pos_weight=pos_weight)
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def adjacency_loss(y_true, y_pred):
    """Implements Smooth-L1 loss.
    y_true and y_pred are typically: [N, 4], but could be any shape.
    """
    print(y_true.shape)
    print(y_pred.shape)
    y_true = tf.cast(y_true, tf.int32)
    total_count = tf.cast(tf.size(y_true), tf.float32)
    count_ones = tf.cast(tf.math.count_nonzero(y_true),
                         tf.float32) + tf.constant([1.0], tf.float32)
    count_zeros = total_count - count_ones
    pos_weight = tf.math.divide(count_zeros, count_ones)
    #loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=tf.one_hot(y_true, depth=2), logits=y_pred)
    loss = tf.nn.weighted_cross_entropy_with_logits(logits=y_pred,
                                                    targets=tf.one_hot(
                                                        y_true, depth=2),
                                                    pos_weight=pos_weight)
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def _build_label(self,
                   image,
                   gt_boxes,
                   gt_classes,
                   info,
                   inds,
                   data,
                   is_training=True):
    """Label construction for both the train and eval data."""
    width = self._image_w
    height = self._image_h

    # Set the image shape.
    imshape = image.get_shape().as_list()
    imshape[-1] = 3
    image.set_shape(imshape)

    labels = dict()
    (labels['inds'], labels['upds'],
     labels['true_conf']) = self._label_builder(gt_boxes, gt_classes, width,
                                                height)

    # Set/fix the boxes shape.
    boxes = self.set_shape(gt_boxes, pad_axis=0, pad_value=0)
    classes = self.set_shape(gt_classes, pad_axis=0, pad_value=-1)

    # Build the dictionary set.
    labels.update({
        'source_id': utils.process_source_id(data['source_id']),
        'bbox': t.cast(boxes, dtype=self._dtype),
        'classes': t.cast(classes, dtype=self._dtype),
        # For OTA loss.
        'image_info': info,
    })

    # Update the labels dictionary.
    if not is_training:
      # Sets up groundtruth data for evaluation.
      groundtruths = {
          'source_id': labels['source_id'],
          'height': data['height'],
          'width': data['width'],
          'num_detections': tf.shape(data['groundtruth_boxes'])[0],
          'image_info': info,
          'boxes': bbox_ops.denormalize_boxes(
              data['groundtruth_boxes'],
              tf.cast([data['height'], data['width']], gt_boxes.dtype)),
          'classes': data['groundtruth_classes'],
          'areas': data['groundtruth_area'],
          'is_crowds': tf.cast(
              tf.gather(data['groundtruth_is_crowd'], inds), tf.int32),
      }
      groundtruths['source_id'] = utils.process_source_id(
          groundtruths['source_id'])
      groundtruths = utils.pad_groundtruths_to_fixed_size(
          groundtruths, self._max_num_instances)
      labels['groundtruths'] = groundtruths
    return image, labels
//  [python]
def _build_label(self,
                   image,
                   gt_boxes,
                   gt_classes,
                   info,
                   inds,
                   data,
                   is_training=True):
    """Label construction for both the train and eval data."""
    width = self._image_w
    height = self._image_h

    # Set the image shape.
    imshape = image.get_shape().as_list()
    imshape[-1] = 3
    image.set_shape(imshape)

    labels = dict()
    (labels['inds'], labels['upds'],
     labels['true_conf']) = self._label_builder(gt_boxes, gt_classes, width,
                                                height)

    # Set/fix the boxes shape.
    boxes = self.set_shape(gt_boxes, pad_axis=0, pad_value=0)
    classes = self.set_shape(gt_classes, pad_axis=0, pad_value=-1)

    # Build the dictionary set.
    labels.update({
        'source_id': utils.process_source_id(data['source_id']),
        'bbox': tf.cast(boxes, dtype=self._dtype),
        'classes': tf.cast(classes, dtype=self._dtype),
        # For OTA loss.
        'image_info': info,
    })

    # Update the labels dictionary.
    if not is_training:
      # Sets up groundtruth data for evaluation.
      groundtruths = {
          'source_id': labels['source_id'],
          'height': data['height'],
          'width': data['width'],
          'num_detections': tf.shape(data['groundtruth_boxes'])[0],
          'image_info': info,
          'boxes': bbox_ops.denormalize_boxes(
              data['groundtruth_boxes'],
              tf.cast([data['height'], data['width']], gt_boxes.dtype)),
          'classes': data['groundtruth_classes'],
          'areas': data['groundtruth_area'],
          'is_crowds': tf.cast(
              tf.gather(data['groundtruth_is_crowd'], inds), tf.int32),
      }
      groundtruths['source_id'] = utils.process_source_id(
          groundtruths['source_id'])
      groundtruths = utils.pad_groundtruths_to_fixed_size(
          groundtruths, self._max_num_instances)
      labels['groundtruths'] = groundtruths
    return image, labels
//  [python]
def trim_zeros_graph(boxes, name='trim_zeros'):
    """Often boxes are represented with matrices of shape [N, 4] and
    are padded with zeros. This removes zero boxes.

    boxes: [N, 4] matrix of boxes.
    non_zeros: [N] a 1D boolean mask identifying the rows to keep
    """
    non_zeros = tf.cast(tf.reduce_sum(tf.abs(boxes), axis=1), tf.bool)
    boxes = tf.boolean_mask(boxes, non_zeros, name=name)
    return boxes, non_zeros
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.

    target_bbox: [batch, num_rois, (dy, dx, log(dh), log(dw))]
    target_class_ids: [batch, num_rois]. Integer class IDs.
    pred_bbox: [batch, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    """
    # Reshape to merge batch and roi dimensions for simplicity.
    target_class_ids = K.reshape(target_class_ids, (-1, ))
    target_bbox = K.reshape(target_bbox, (-1, 4))
    pred_bbox = K.reshape(pred_bbox, (-1, K.int_shape(pred_bbox)[2], 4))

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indices.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_roi_ix))
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_bbox = tf.gather(target_bbox, positive_roi_ix)
    pred_bbox = tf.gather_nd(pred_bbox, indices)

    # Smooth-L1 Loss
    loss = K.switch(
        tf.size(target_bbox) > 0,
        smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox), tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.

    target_bbox: [batch, num_rois, (dy, dx, log(dh), log(dw))]
    target_class_ids: [batch, num_rois]. Integer class IDs.
    pred_bbox: [batch, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    """
    # Reshape to merge batch and roi dimensions for simplicity.
    target_class_ids = K.reshape(target_class_ids, (-1, ))
    target_bbox = K.reshape(target_bbox, (-1, 4))
    pred_bbox = K.reshape(pred_bbox, (-1, K.int_shape(pred_bbox)[2], 4))

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indices.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_roi_ix), tf.int64)
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_bbox = tf.gather(target_bbox, positive_roi_ix)
    pred_bbox = tf.gather_nd(pred_bbox, indices)

    # Smooth-L1 Loss
    loss = K.switch(
        tf.size(target_bbox) > 0,
        smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox), tf.constant(0.0))
    loss = K.mean(loss)
    return loss