jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def prototypical_clustering_learn_layer(nclasses,
                                        x_train,
                                        y_train,
                                        x_test,
                                        phi,
                                        num_cluster_steps,
                                        lambd=0.1,
                                        alpha=0.0):
  """Computes the prototypes, cluster centers, with additional clustering on
    the validation data.
  Args:
    x_train: [N, D], Train data.
    y_train: [N], Train class labels.
    x_test: [N, D], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, wts = phi(x_all, reuse=None, is_training=True)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]
  wts_keys = wts.keys()
  wts_tensors = [wts[kk] for kk in wts_keys]

  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.to_float(tf.equal(y_train, kk)),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)

  protos = concat(protos, 0)  # [K, D]

  for tt in range(num_cluster_steps):
    # Compute new representation of the data.
    h, _ = phi(
        x_all,
        reuse=True,
        is_training=True,
        ext_wts=dict(zip(wts_keys, wts_tensors)))
    h_train = h[:num_x_train, :]
    h_test = h[num_x_train:, :]
    all_data = concat([h_train, h_test], 0)

    # Label assignment.
    prob = assign_cluster(protos, all_data)

    # Prototype update.
    ## This is a vanilla version.
    ## Probably want to impose a constraint that each training example can only
    ## be in one cluster.
    protos = update_cluster(all_data, prob)

    if alpha > 0.0 and tt < num_cluster_steps - 1:
      # We can also use soft labels here.
      loss = lambd * sq_dist_loss(protos, all_data)
      # One gradient update towards fast weights.
      print(wts_tensors)
      [print(vv.name) for vv in wts_tensors]
      grads = tf.gradients(loss, wts_tensors, gate_gradients=True)
      [print(gg) for gg in grads]
      # Stop the gradient of the gradient.
      grads = [tf.stop_gradient(gg) for gg in grads]
      wts_tensors = [wt - alpha * gg for wt, gg in zip(wts_tensors, grads)]

  # Be cautious here!!
  #uloss = sq_dist_loss(protos, all_data)
  uloss = 0.0
  logits = compute_logits(protos, h_test)
  return logits, uloss
//  [python]
def prototypical_clustering_learn_layer(nclasses,
                                        x_train,
                                        y_train,
                                        x_test,
                                        phi,
                                        num_cluster_steps,
                                        lambd=0.1,
                                        alpha=0.0):
  """Computes the prototypes, cluster centers, with additional clustering on
    the validation data.
  Args:
    x_train: [N, D], Train data.
    y_train: [N], Train class labels.
    x_test: [N, D], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, wts = phi(x_all, reuse=None, is_training=True)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]
  wts_keys = wts.keys()
  wts_tensors = [wts[kk] for kk in wts_keys]

  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.cast(tf.equal(y_train, kk), h_train.dtype),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)

  protos = concat(protos, 0)  # [K, D]

  for tt in range(num_cluster_steps):
    # Compute new representation of the data.
    h, _ = phi(
        x_all,
        reuse=True,
        is_training=True,
        ext_wts=dict(zip(wts_keys, wts_tensors)))
    h_train = h[:num_x_train, :]
    h_test = h[num_x_train:, :]
    all_data = concat([h_train, h_test], 0)

    # Label assignment.
    prob = assign_cluster(protos, all_data)

    # Prototype update.
    ## This is a vanilla version.
    ## Probably want to impose a constraint that each training example can only
    ## be in one cluster.
    protos = update_cluster(all_data, prob)

    if alpha > 0.0 and tt < num_cluster_steps - 1:
      # We can also use soft labels here.
      loss = lambd * sq_dist_loss(protos, all_data)
      # One gradient update towards fast weights.
      print(wts_tensors)
      [print(vv.name) for vv in wts_tensors]
      grads = tf.gradients(loss, wts_tensors, gate_gradients=True)
      [print(gg) for gg in grads]
      # Stop the gradient of the gradient.
      grads = [tf.stop_gradient(gg) for gg in grads]
      wts_tensors = [wt - alpha * gg for wt, gg in zip(wts_tensors, grads)]

  # Be cautious here!!
  #uloss = sq_dist_loss(protos, all_data)
  uloss = 0.0
  logits = compute_logits(protos, h_test)
  return logits, uloss
//  [python]
def _bmcalc(self, y):
        """
        Calculate branch metrics for a given noisy codeword tensor.
        For each time period t, _bmcalc computes the distance of symbol
        vector y[t] from each possible output symbol.
        The distance metric is L2 distance if decoder parameter `method` is
        "soft".

        The distance metric is L1 distance if parameter `method` is "hard".
        """

        op_bits = np.stack(
                    [int2bin(op, self._conv_n) for op in range(self._no)])
        op_mat = t.cast(t.tile(op_bits, [1,self._num_syms]), t.float32)
        op_mat = t.expand_dims(op_mat, axis=0)
        y = t.expand_dims(y, axis=1)
        if self._method=='soft_llr':
            op_mat_sign = 1 - 2.*op_mat
            llr_sign = -1. * t.math.multiply(y, op_mat_sign)
            llr_sign = t.reshape(llr_sign,
                (-1, self._no, self._num_syms, self._conv_n))
            # Sum of LLR*(sign of bit) for each symbol
            bm = t.math.reduce_sum(llr_sign, axis=-1)

        elif self._method == 'hard':
            diffabs = t.math.abs(y-op_mat)
            diffabs = t.reshape(diffabs,
                                 (-1, self._no, self._num_syms, self._conv_n))
            # Manhattan distance of symbols
            bm = t.math.reduce_sum(diffabs, axis=-1)

        return bm
//  [python]
def _bmcalc(self, y):
        """
        Calculate branch metrics for a given noisy codeword tensor.
        For each time period t, _bmcalc computes the distance of symbol
        vector y[t] from each possible output symbol.
        The distance metric is L2 distance if decoder parameter `method` is
        "soft".

        The distance metric is L1 distance if parameter `method` is "hard".
        """

        op_bits = np.stack(
                    [int2bin(op, self._conv_n) for op in range(self._no)])
        op_mat = tf.cast(tf.tile(op_bits, [1,self._num_syms]), tf.float32)
        op_mat = tf.expand_dims(op_mat, axis=0)
        y = tf.expand_dims(y, axis=1)
        if self._method=='soft_llr':
            op_mat_sign = 1 - 2.*op_mat
            llr_sign = -1. * tf.math.multiply(y, op_mat_sign)
            llr_sign = tf.reshape(llr_sign,
                (-1, self._no, self._num_syms, self._conv_n))
            # Sum of LLR*(sign of bit) for each symbol
            bm = tf.math.reduce_sum(llr_sign, axis=-1)

        elif self._method == 'hard':
            diffabs = tf.math.abs(y-op_mat)
            diffabs = tf.reshape(diffabs,
                                 (-1, self._no, self._num_syms, self._conv_n))
            # Manhattan distance of symbols
            bm = tf.math.reduce_sum(diffabs, axis=-1)

        return bm
//  [python]
def _bmcalc(self, llr_in):
        """
        Calculate branch gamma metrics for a given noisy codeword tensor.
        For each time period t, _bmcalc computes the "distance" of symbol
        vector y[t] from each possible output symbol i.e.,
        (2*Eb/N0)* sum_i x_y*y_i for i=1,2,...,conv_n

        The above metric is used in calculation of gamma.
        If the input is llr, which is nothing but 2*Eb*y/N0.
        """
        op_bits = np.stack(
                    [int2bin(op, self._conv_n) for op in range(self._no)])
        op_mat = t.cast(t.tile(op_bits, [1, self._num_syms]), t.float32)
        op_mat = t.expand_dims(op_mat, axis=0)
        llr_in = t.expand_dims(llr_in, axis=1)
        op_mat_sign = 1. - 2. * op_mat

        llr_sign = t.math.multiply(llr_in, op_mat_sign)
        half_llr_sign = t.reshape(0.5 * llr_sign,
            (-1, self._no, self._num_syms, self._conv_n))

        if self._algorithm in ['log', 'maxlog']:
            bm = t.math.reduce_sum(half_llr_sign, axis=-1)
        else:
            bm = t.math.exp(t.math.reduce_sum(half_llr_sign, axis=-1))

        return bm
//  [python]
def _bmcalc(self, llr_in):
        """
        Calculate branch gamma metrics for a given noisy codeword tensor.
        For each time period t, _bmcalc computes the "distance" of symbol
        vector y[t] from each possible output symbol i.e.,
        (2*Eb/N0)* sum_i x_y*y_i for i=1,2,...,conv_n

        The above metric is used in calculation of gamma.
        If the input is llr, which is nothing but 2*Eb*y/N0.
        """
        op_bits = np.stack(
                    [int2bin(op, self._conv_n) for op in range(self._no)])
        op_mat = tf.cast(tf.tile(op_bits, [1, self._num_syms]), tf.float32)
        op_mat = tf.expand_dims(op_mat, axis=0)
        llr_in = tf.expand_dims(llr_in, axis=1)
        op_mat_sign = 1. - 2. * op_mat

        llr_sign = tf.math.multiply(llr_in, op_mat_sign)
        half_llr_sign = tf.reshape(0.5 * llr_sign,
            (-1, self._no, self._num_syms, self._conv_n))

        if self._algorithm in ['log', 'maxlog']:
            bm = tf.math.reduce_sum(half_llr_sign, axis=-1)
        else:
            bm = tf.math.exp(tf.math.reduce_sum(half_llr_sign, axis=-1))

        return bm
//  [python]
def call(self, inputs):
        """
        Viterbi decoding function.

        inputs is the (noisy) codeword tensor where the last dimension should
        equal n. All the leading dimensions are assumed as batch dimensions.

        """
        LARGEDIST = 2.**20 # pylint: disable=invalid-name

        tf.debugging.assert_type(inputs, tf.float32,
                                 message="input must be tf.float32.")
        if self._method == 'hard':
            inputs = tf.math.floormod(t.cast(inputs, tf.int32),2)
        elif self._method == 'soft_llr':
            inputs = -1. * inputs

        inputs = t.cast(inputs, tf.float32)

        output_shape = inputs.get_shape().as_list()
        y_resh = tf.reshape(inputs, [-1, self._n])
        output_shape[0] = -1
        if self._return_info_bits:
            output_shape[-1] = self._k # assign k to the last dimension
        else:
            output_shape[-1] = self._n
        # Branch metrics matrix for a given y
        bm_mat = self._bmcalc(y_resh)

        init_cm_np = np.full((self._ns,), LARGEDIST)
        init_cm_np[0] = 0.0
        prev_cm_ = tf.convert_to_tensor(init_cm_np, dtype=tf.float32)
        prev_cm = tf.tile(prev_cm_[None,:], [tf.shape(y_resh)[0], 1])

        cm_ta, tb_ta = self._update_fwd(prev_cm, bm_mat)

        cm = tf.transpose(cm_ta.stack(), perm=[1,2,0])
        tb = tf.transpose(tb_ta.stack(),perm=[1,2,0])
        del cm_ta, tb_ta

        zero_st = tf.zeros((tf.shape(y_resh)[0], 1), tf.int32)
        opt_path = self._optimal_path(cm, tb)
        opt_path = tf.concat((zero_st, opt_path), axis=1)
        del cm, tb
        msghat, cwhat = self._op_bits_path(opt_path)
        if self._return_info_bits:
            msghat = msghat[...,:self._k]
            output = t.cast(msghat, self.output_dtype)
        else:
            output = t.cast(cwhat, self.output_dtype)
        output_reshaped = tf.reshape(output, output_shape)

        return output_reshaped
//  [python]
def call(self, inputs):
        """
        Viterbi decoding function.

        inputs is the (noisy) codeword tensor where the last dimension should
        equal n. All the leading dimensions are assumed as batch dimensions.

        """
        LARGEDIST = 2.**20 # pylint: disable=invalid-name

        tf.debugging.assert_type(inputs, tf.float32,
                                 message="input must be tf.float32.")
        if self._method == 'hard':
            inputs = tf.math.floormod(tf.cast(inputs, tf.int32),2)
        elif self._method == 'soft_llr':
            inputs = -1. * inputs

        inputs = tf.cast(inputs, tf.float32)

        output_shape = inputs.get_shape().as_list()
        y_resh = tf.reshape(inputs, [-1, self._n])
        output_shape[0] = -1
        if self._return_info_bits:
            output_shape[-1] = self._k # assign k to the last dimension
        else:
            output_shape[-1] = self._n
        # Branch metrics matrix for a given y
        bm_mat = self._bmcalc(y_resh)

        init_cm_np = np.full((self._ns,), LARGEDIST)
        init_cm_np[0] = 0.0
        prev_cm_ = tf.convert_to_tensor(init_cm_np, dtype=tf.float32)
        prev_cm = tf.tile(prev_cm_[None,:], [tf.shape(y_resh)[0], 1])

        cm_ta, tb_ta = self._update_fwd(prev_cm, bm_mat)

        cm = tf.transpose(cm_ta.stack(), perm=[1,2,0])
        tb = tf.transpose(tb_ta.stack(),perm=[1,2,0])
        del cm_ta, tb_ta

        zero_st = tf.zeros((tf.shape(y_resh)[0], 1), tf.int32)
        opt_path = self._optimal_path(cm, tb)
        opt_path = tf.concat((zero_st, opt_path), axis=1)
        del cm, tb
        msghat, cwhat = self._op_bits_path(opt_path)
        if self._return_info_bits:
            msghat = msghat[...,:self._k]
            output = tf.cast(msghat, self.output_dtype)
        else:
            output = tf.cast(cwhat, self.output_dtype)
        output_reshaped = tf.reshape(output, output_shape)

        return output_reshaped
//  [python]
def get_model_rollout(uvf_agent, tf_env):
  """Model rollout function."""
  state_spec = tf_env.observation_spec()[0]
  action_spec = tf_env.action_spec()[0]
  state_ph = tf.placeholder(dtype=state_spec.dtype, shape=state_spec.shape)
  action_ph = tf.placeholder(dtype=action_spec.dtype, shape=action_spec.shape)

  merged_state = uvf_agent.merged_state(state_ph)
  diff_value = uvf_agent.critic_net(tf.expand_dims(merged_state, 0),
                                    tf.expand_dims(action_ph, 0))[0]
  diff_value = tf.to_float(diff_value)
  state_ph.shape.assert_is_compatible_with(diff_value.shape)
  next_state = state_ph + diff_value

  def model_rollout_fn(sess, state, action):
    return sess.run(next_state, feed_dict={state_ph: state, action_ph: action})

  return model_rollout_fn
//  [python]
def get_model_rollout(uvf_agent, tf_env):
  """Model rollout function."""
  state_spec = tf_env.observation_spec()[0]
  action_spec = tf_env.action_spec()[0]
  state_ph = tf.placeholder(dtype=state_spec.dtype, shape=state_spec.shape)
  action_ph = tf.placeholder(dtype=action_spec.dtype, shape=action_spec.shape)

  merged_state = uvf_agent.merged_state(state_ph)
  diff_value = uvf_agent.critic_net(tf.expand_dims(merged_state, 0),
                                    tf.expand_dims(action_ph, 0))[0]
  diff_value = tf.cast(diff_value, dtype=state_ph.dtype)
  state_ph.shape.assert_is_compatible_with(diff_value.shape)
  next_state = state_ph + diff_value

  def model_rollout_fn(sess, state, action):
    return sess.run(next_state, feed_dict={state_ph: state, action_ph: action})

  return model_rollout_fn
//  [python]
def keypoints_absolute2relative(keypoints, width, height):
    if isinstance(height,tf.Tensor) and height.dtype != keypoints.dtype:
        height = t.cast(height, keypoints.dtype)
    if isinstance(width,tf.Tensor) and width.dtype != keypoints.dtype:
        width = t.cast(width, keypoints.dtype)
    x = keypoints[...,0] / (width-1)
    y = keypoints[...,1] / (height-1)

    x_in_range = tf.logical_and(tf.greater_equal(x,0),tf.less_equal(x,1.001))
    x = tf.where(x_in_range,x,-tf.ones_like(x))
    y_in_range = tf.logical_and(tf.greater_equal(y,0),tf.less_equal(y,1.001))
    y = tf.where(y_in_range,y,-tf.ones_like(y))

    return tf.stack([x,y], axis=-1)
//  [python]
def keypoints_absolute2relative(keypoints, width, height):
    if isinstance(height,tf.Tensor) and height.dtype != keypoints.dtype:
        height = tf.cast(height, keypoints.dtype)
    if isinstance(width,tf.Tensor) and width.dtype != keypoints.dtype:
        width = tf.cast(width, keypoints.dtype)
    x = keypoints[...,0] / (width-1)
    y = keypoints[...,1] / (height-1)

    x_in_range = tf.logical_and(tf.greater_equal(x,0),tf.less_equal(x,1.001))
    x = tf.where(x_in_range,x,-tf.ones_like(x))
    y_in_range = tf.logical_and(tf.greater_equal(y,0),tf.less_equal(y,1.001))
    y = tf.where(y_in_range,y,-tf.ones_like(y))

    return tf.stack([x,y], axis=-1)
//  [python]
def keypoints_relative2absolute(keypoints,width,height):
    if not isinstance(width,tf.Tensor):
        width = tf.convert_to_tensor(width,dtype=keypoints.dtype)
    elif width.dtype != keypoints.dtype:
        width = t.cast(width,keypoints.dtype)
    if not isinstance(height,tf.Tensor):
        height = tf.convert_to_tensor(height,dtype=keypoints.dtype)
    elif height.dtype != keypoints.dtype:
        height = t.cast(height,keypoints.dtype)
    x = keypoints[...,0]*(width-1)
    y = keypoints[...,1]*(height-1)

    x_in_range = tf.logical_and(tf.greater_equal(x,0),tf.less_equal(x,width+2))
    x = tf.where(x_in_range,x,-tf.ones_like(x))
    y_in_range = tf.logical_and(tf.greater_equal(y,0),tf.less_equal(y,height+2))
    y = tf.where(y_in_range,y,-tf.ones_like(y))

    return tf.stack([x,y],axis=-1)
//  [python]
def keypoints_relative2absolute(keypoints,width,height):
    if not isinstance(width,tf.Tensor):
        width = tf.convert_to_tensor(width,dtype=keypoints.dtype)
    elif width.dtype != keypoints.dtype:
        width = tf.cast(width,keypoints.dtype)
    if not isinstance(height,tf.Tensor):
        height = tf.convert_to_tensor(height,dtype=keypoints.dtype)
    elif height.dtype != keypoints.dtype:
        height = tf.cast(height,keypoints.dtype)
    x = keypoints[...,0]*(width-1)
    y = keypoints[...,1]*(height-1)

    x_in_range = tf.logical_and(tf.greater_equal(x,0),tf.less_equal(x,width+2))
    x = tf.where(x_in_range,x,-tf.ones_like(x))
    y_in_range = tf.logical_and(tf.greater_equal(y,0),tf.less_equal(y,height+2))
    y = tf.where(y_in_range,y,-tf.ones_like(y))

    return tf.stack([x,y],axis=-1)
//  [python]
def _build_inputs(self, image):
    """Builds classification model inputs for serving."""
    if isinstance(image, tf.RaggedTensor):
      image = image.to_tensor()
    image = t.cast(image, dtype=tf.float32)

    # Normalizes image with mean and std pixel values.
    image_feature = self.params.task.train_data.image_feature
    image = preprocess_ops.normalize_image(
        image,
        offset=image_feature.mean,
        scale=image_feature.stddev)

    if self.params.task.train_data.preserve_aspect_ratio:
      image, image_info = preprocess_ops.resize_and_crop_image(
          image,
          self._input_image_size,
          padded_size=self._input_image_size,
          aug_scale_min=1.0,
          aug_scale_max=1.0)
    else:
      image, image_info = preprocess_ops.resize_image(image,
                                                      self._input_image_size)
    return image, image_info
//  [python]
def _build_inputs(self, image):
    """Builds classification model inputs for serving."""
    if isinstance(image, tf.RaggedTensor):
      image = image.to_tensor()
    image = tf.cast(image, dtype=tf.float32)

    # Normalizes image with mean and std pixel values.
    image_feature = self.params.task.train_data.image_feature
    image = preprocess_ops.normalize_image(
        image,
        offset=image_feature.mean,
        scale=image_feature.stddev)

    if self.params.task.train_data.preserve_aspect_ratio:
      image, image_info = preprocess_ops.resize_and_crop_image(
          image,
          self._input_image_size,
          padded_size=self._input_image_size,
          aug_scale_min=1.0,
          aug_scale_max=1.0)
    else:
      image, image_info = preprocess_ops.resize_image(image,
                                                      self._input_image_size)
    return image, image_info
//  [python]
def resize_fixed_percentage(tensors, size):
  assert len(size) == 1
  percentage = float(size[0])/100.0
  h = t.cast(t.shape(tensors[DataKeys.IMAGES])[0], t.float32)
  w = t.cast(t.shape(tensors[DataKeys.IMAGES])[1], t.float32)

  newh = t.multiply(percentage, h)
  neww = t.multiply(percentage, w)

  newh = t.cast(t.round(newh), t.int32)
  neww = t.cast(t.round(neww), t.int32)

  return resize_fixed_size(tensors, t.stack([newh, neww], axis=0))
//  [python]
def resize_fixed_percentage(tensors, size):
  assert len(size) == 1
  percentage = float(size[0])/100.0
  h = tf.cast(tf.shape(tensors[DataKeys.IMAGES])[0], tf.float32)
  w = tf.cast(tf.shape(tensors[DataKeys.IMAGES])[1], tf.float32)

  newh = tf.multiply(percentage, h)
  neww = tf.multiply(percentage, w)

  newh = tf.cast(tf.round(newh), tf.int32)
  neww = tf.cast(tf.round(neww), tf.int32)

  return resize_fixed_size(tensors, tf.stack([newh, neww], axis=0))
//  [python]
def _update_fwd(self, alph_init, bm_mat, llr):
        """
        Run forward update from time t=0 to t=k-1.
        At each time t, computes alpha_t using alpha_t-1 and gamma_t.

        Returns tensor array of alpha_t, t-0,1,2...,k-1
        """
        alph_ta = tf.TensorArray(tf.float32, size=self._num_syms+1,
                                  dynamic_size=False, clear_after_read=False)
        alph_prev = t.cast(alph_init, tf.float32)

        # (bs, _Ns, _ni, 2) matrix
        ipst_ip_mask = tf.tile(
            self.ipst_ip_idx[None,:],[tf.shape(alph_init)[0],1,1,1])
        # (bs, _Ns, _ni) matrix, by from state
        op_mask = tf.tile(self.trellis.op_by_fromnode[None,:,:],
                          [tf.shape(alph_init)[0],1,1])
        ipbit_mat = tf.tile(tf.range(self._ni)[None, None, :],
                            [tf.shape(alph_init)[0], self._ns, 1])
        ipbitsign_mat = 1. - 2. * t.cast(ipbit_mat, tf.float32)
        alph_ta = alph_ta.write(0, alph_prev)
        for t in tf.range(self._num_syms):
            bm_t = bm_mat[..., t]
            llr_t = 0.5 * llr[...,t][:, None,None]

            bm_byfromst = tf.gather(bm_t, op_mask, batch_dims=1)
            signed_half_llr = tf.math.multiply(
                tf.tile(llr_t,[1, self._ns, self._ni]), ipbitsign_mat)
            if self._algorithm in ['log', 'maxlog']:
                llr_byfromst = signed_half_llr
                gamma_byfromst = llr_byfromst + bm_byfromst
                alph_gam_prod = gamma_byfromst + alph_prev[:,:,None]
            else:
                llr_byfromst = tf.math.exp(signed_half_llr)
                gamma_byfromst = tf.multiply(llr_byfromst, bm_byfromst)
                alph_gam_prod = tf.math.multiply(gamma_byfromst,
                                             alph_prev[:,:,None])

            alphgam_bytost = tf.gather_nd(alph_gam_prod,
                                          ipst_ip_mask,
                                          batch_dims=1)
            if self._algorithm =='map':
                alph_t = tf.math.reduce_sum(alphgam_bytost, axis=-1)
                alph_t_sum = tf.reduce_sum(alph_t, axis=-1)
                alph_t = tf.divide(alph_t, tf.tile(alph_t_sum[:,None],[1,self._ns]))
            elif self._algorithm == 'log':
                alph_t = tf.math.reduce_logsumexp(alphgam_bytost, axis=-1)
            else:  # self._algorithm = 'maxlog'
                alph_t = tf.math.reduce_max(alphgam_bytost, axis=-1)

            alph_prev = alph_t
            alph_ta = alph_ta.write(t+1, alph_t)
        return alph_ta
//  [python]
def _update_fwd(self, alph_init, bm_mat, llr):
        """
        Run forward update from time t=0 to t=k-1.
        At each time t, computes alpha_t using alpha_t-1 and gamma_t.

        Returns tensor array of alpha_t, t-0,1,2...,k-1
        """
        alph_ta = tf.TensorArray(tf.float32, size=self._num_syms+1,
                                  dynamic_size=False, clear_after_read=False)
        alph_prev = tf.cast(alph_init, tf.float32)

        # (bs, _Ns, _ni, 2) matrix
        ipst_ip_mask = tf.tile(
            self.ipst_ip_idx[None,:],[tf.shape(alph_init)[0],1,1,1])
        # (bs, _Ns, _ni) matrix, by from state
        op_mask = tf.tile(self.trellis.op_by_fromnode[None,:,:],
                          [tf.shape(alph_init)[0],1,1])
        ipbit_mat = tf.tile(tf.range(self._ni)[None, None, :],
                            [tf.shape(alph_init)[0], self._ns, 1])
        ipbitsign_mat = 1. - 2. * tf.cast(ipbit_mat, tf.float32)
        alph_ta = alph_ta.write(0, alph_prev)
        for t in tf.range(self._num_syms):
            bm_t = bm_mat[..., t]
            llr_t = 0.5 * llr[...,t][:, None,None]

            bm_byfromst = tf.gather(bm_t, op_mask, batch_dims=1)
            signed_half_llr = tf.math.multiply(
                tf.tile(llr_t,[1, self._ns, self._ni]), ipbitsign_mat)
            if self._algorithm in ['log', 'maxlog']:
                llr_byfromst = signed_half_llr
                gamma_byfromst = llr_byfromst + bm_byfromst
                alph_gam_prod = gamma_byfromst + alph_prev[:,:,None]
            else:
                llr_byfromst = tf.math.exp(signed_half_llr)
                gamma_byfromst = tf.multiply(llr_byfromst, bm_byfromst)
                alph_gam_prod = tf.math.multiply(gamma_byfromst,
                                             alph_prev[:,:,None])

            alphgam_bytost = tf.gather_nd(alph_gam_prod,
                                          ipst_ip_mask,
                                          batch_dims=1)
            if self._algorithm =='map':
                alph_t = tf.math.reduce_sum(alphgam_bytost, axis=-1)
                alph_t_sum = tf.reduce_sum(alph_t, axis=-1)
                alph_t = tf.divide(alph_t, tf.tile(alph_t_sum[:,None],[1,self._ns]))
            elif self._algorithm == 'log':
                alph_t = tf.math.reduce_logsumexp(alphgam_bytost, axis=-1)
            else:  # self._algorithm = 'maxlog'
                alph_t = tf.math.reduce_max(alphgam_bytost, axis=-1)

            alph_prev = alph_t
            alph_ta = alph_ta.write(t+1, alph_t)
        return alph_ta
//  [python]
def call(self, inputs: tf.Tensor):
        """Calls `text.BertTokenizer` on inputs.

        Args:
          inputs: A string Tensor of shape `(batch_size,)`.

        Returns:
          One or three of `RaggedTensors` if `tokenize_with_offsets` is False or
          True, respectively. These are
            tokens: A `RaggedTensor` of shape
              `[batch_size, (words), (pieces_per_word)]`
              and type int32. `tokens[i,j,k]` contains the k-th wordpiece of the
              j-th word in the i-th input.
            start_offsets, limit_offsets: If `tokenize_with_offsets` is True,
              RaggedTensors of type int64 with the same indices as tokens.
              Element `[i,j,k]` contains the byte offset at the start, or past the
              end, resp., for the k-th wordpiece of the j-th word in the i-th input.
        """
        # Prepare to reshape the result to work around broken shape inference.
        batch_size = tf.shape(inputs)[0]

        def _reshape(rt):
            values = rt.values
            row_splits = rt.row_splits
            row_splits = tf.reshape(row_splits, [batch_size + 1])
            return tf.RaggedTensor.from_row_splits(values, row_splits)

        # Call the tokenizer.
        if self.tokenize_with_offsets:
            tokens, start_offsets, limit_offsets = self._bert_tokenizer.tokenize_with_offsets(inputs)
            tokens = t.cast(tokens, dtype=tf.int32)
            return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
        else:
            tokens = self._bert_tokenizer.tokenize(inputs)
            tokens = t.cast(tokens, dtype=tf.int32)
            return _reshape(tokens)
//  [python]
def call(self, inputs: tf.Tensor):
        """Calls `text.BertTokenizer` on inputs.

        Args:
          inputs: A string Tensor of shape `(batch_size,)`.

        Returns:
          One or three of `RaggedTensors` if `tokenize_with_offsets` is False or
          True, respectively. These are
            tokens: A `RaggedTensor` of shape
              `[batch_size, (words), (pieces_per_word)]`
              and type int32. `tokens[i,j,k]` contains the k-th wordpiece of the
              j-th word in the i-th input.
            start_offsets, limit_offsets: If `tokenize_with_offsets` is True,
              RaggedTensors of type int64 with the same indices as tokens.
              Element `[i,j,k]` contains the byte offset at the start, or past the
              end, resp., for the k-th wordpiece of the j-th word in the i-th input.
        """
        # Prepare to reshape the result to work around broken shape inference.
        batch_size = tf.shape(inputs)[0]

        def _reshape(rt):
            values = rt.values
            row_splits = rt.row_splits
            row_splits = tf.reshape(row_splits, [batch_size + 1])
            return tf.RaggedTensor.from_row_splits(values, row_splits)

        # Call the tokenizer.
        if self.tokenize_with_offsets:
            tokens, start_offsets, limit_offsets = self._bert_tokenizer.tokenize_with_offsets(inputs)
            tokens = tf.cast(tokens, dtype=tf.int32)
            return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
        else:
            tokens = self._bert_tokenizer.tokenize(inputs)
            tokens = tf.cast(tokens, dtype=tf.int32)
            return _reshape(tokens)
//  [python]
def __init__(self,
               config: Optional[ModelConfig] = None,
               overrides: Optional[Dict[Text, Any]] = None):
    """Create an EfficientNet model.

    Args:
      config: (optional) the main model parameters to create the model
      overrides: (optional) a dict containing keys that can override config
    """
    overrides = overrides or {}
    config = config or ModelConfig()

    self.config = config.replace(**overrides)

    input_channels = self.config.input_channels
    model_name = self.config.model_name
    input_shape = (None, None, input_channels)  # Should handle any size image
    image_input = tf_keras.layers.Input(shape=input_shape)

    output = efficientnet(image_input, self.config)

    # Cast to float32 in case we have a different model dtype
    output = tf.to_float(output)

    logging.info('Building model %s with params %s', model_name, self.config)

    super(EfficientNet, self).__init__(
        inputs=image_input, outputs=output, name=model_name)
//  [python]
def __init__(self,
               config: Optional[ModelConfig] = None,
               overrides: Optional[Dict[Text, Any]] = None):
    """Create an EfficientNet model.

    Args:
      config: (optional) the main model parameters to create the model
      overrides: (optional) a dict containing keys that can override config
    """
    overrides = overrides or {}
    config = config or ModelConfig()

    self.config = config.replace(**overrides)

    input_channels = self.config.input_channels
    model_name = self.config.model_name
    input_shape = (None, None, input_channels)  # Should handle any size image
    image_input = tf_keras.layers.Input(shape=input_shape)

    output = efficientnet(image_input, self.config)

    # Cast to float32 in case we have a different model dtype
    output = tf.cast(output, tf.float32)

    logging.info('Building model %s with params %s', model_name, self.config)

    super(EfficientNet, self).__init__(
        inputs=image_input, outputs=output, name=model_name)
//  [python]
def round_robin_truncate_inputs(
    inputs: Union[tf.RaggedTensor, List[tf.RaggedTensor]],
    limit: Union[int, tf.Tensor],
) -> Union[tf.RaggedTensor, List[tf.RaggedTensor]]:
    """Truncates a list of batched segments to fit a per-example length limit.

    Available space is assigned one token at a time in a round-robin fashion
    to the inputs that still need some, until the limit is reached.
    (Or equivalently: the longest input is truncated by one token until the total
    length of inputs fits the limit.) Examples that fit the limit as passed in
    remain unchanged.

    Args:
      inputs: A list of rank-2 RaggedTensors. The i-th example is given by
        the i-th row in each list element, that is, `inputs[:][i, :]`.
      limit: The largest permissible number of tokens in total across one example.

    Returns:
      A list of rank-2 RaggedTensors at corresponding indices with the inputs,
        in which the rows of each RaggedTensor have been truncated such that
        the total number of tokens in each example does not exceed the `limit`.
    """
    if not isinstance(inputs, (list, tuple)):
        return round_robin_truncate_inputs([inputs], limit)[0]
    limit = t.cast(limit, tf.int64)
    if not all(rt.shape.rank == 2 for rt in inputs):
        raise ValueError("All inputs must have shape [batch_size, (items)]")
    if len(inputs) == 1:
        return [_truncate_row_lengths(inputs[0], limit)]
    elif len(inputs) == 2:
        size_a, size_b = [rt.row_lengths() for rt in inputs]
        # Here's a brain-twister: This does round-robin assignment of quota
        # to both inputs until the limit is reached. Hint: consider separately
        # the cases of zero, one, or two inputs exceeding half the limit.
        floor_half = limit // 2
        ceil_half = limit - floor_half
        quota_a = tf.minimum(size_a, ceil_half + tf.nn.relu(floor_half - size_b))
        quota_b = tf.minimum(size_b, floor_half + tf.nn.relu(ceil_half - size_a))
        return [_truncate_row_lengths(inputs[0], quota_a), _truncate_row_lengths(inputs[1], quota_b)]
    else:
        # Note that we don't merge with the 2 input case because the full algorithm
        # is more expensive.
        capacity = tf.stack([rt.row_lengths() for rt in inputs])  # #Segments x B
        new_capacity = _iterative_vectorized_fair_share(capacity, limit)
        return [_truncate_row_lengths(inputs[i], new_capacity[i]) for i in range(capacity.shape[0])]
//  [python]
def round_robin_truncate_inputs(
    inputs: Union[tf.RaggedTensor, List[tf.RaggedTensor]],
    limit: Union[int, tf.Tensor],
) -> Union[tf.RaggedTensor, List[tf.RaggedTensor]]:
    """Truncates a list of batched segments to fit a per-example length limit.

    Available space is assigned one token at a time in a round-robin fashion
    to the inputs that still need some, until the limit is reached.
    (Or equivalently: the longest input is truncated by one token until the total
    length of inputs fits the limit.) Examples that fit the limit as passed in
    remain unchanged.

    Args:
      inputs: A list of rank-2 RaggedTensors. The i-th example is given by
        the i-th row in each list element, that is, `inputs[:][i, :]`.
      limit: The largest permissible number of tokens in total across one example.

    Returns:
      A list of rank-2 RaggedTensors at corresponding indices with the inputs,
        in which the rows of each RaggedTensor have been truncated such that
        the total number of tokens in each example does not exceed the `limit`.
    """
    if not isinstance(inputs, (list, tuple)):
        return round_robin_truncate_inputs([inputs], limit)[0]
    limit = tf.cast(limit, tf.int64)
    if not all(rt.shape.rank == 2 for rt in inputs):
        raise ValueError("All inputs must have shape [batch_size, (items)]")
    if len(inputs) == 1:
        return [_truncate_row_lengths(inputs[0], limit)]
    elif len(inputs) == 2:
        size_a, size_b = [rt.row_lengths() for rt in inputs]
        # Here's a brain-twister: This does round-robin assignment of quota
        # to both inputs until the limit is reached. Hint: consider separately
        # the cases of zero, one, or two inputs exceeding half the limit.
        floor_half = limit // 2
        ceil_half = limit - floor_half
        quota_a = tf.minimum(size_a, ceil_half + tf.nn.relu(floor_half - size_b))
        quota_b = tf.minimum(size_b, floor_half + tf.nn.relu(ceil_half - size_a))
        return [_truncate_row_lengths(inputs[0], quota_a), _truncate_row_lengths(inputs[1], quota_b)]
    else:
        # Note that we don't merge with the 2 input case because the full algorithm
        # is more expensive.
        capacity = tf.stack([rt.row_lengths() for rt in inputs])  # #Segments x B
        new_capacity = _iterative_vectorized_fair_share(capacity, limit)
        return [_truncate_row_lengths(inputs[i], new_capacity[i]) for i in range(capacity.shape[0])]
//  [python]
def object_crop_fixed_offset(tensors, size):
  label = tensors[DataKeys.SEGMENTATION_LABELS]
  object_locations = t.cast(t.where(t.not_equal(label, 0))[:, :2], t.int32)
  min_val = t.maximum(t.constant([0, 0]),
                       t.reduce_max(object_locations, axis=0) - size)
  offset = t.concat([min_val, [0]], axis=0)

  return random_crop_tensors(tensors, size, offset)[0]
//  [python]
def object_crop_fixed_offset(tensors, size):
  label = tensors[DataKeys.SEGMENTATION_LABELS]
  object_locations = tf.cast(tf.where(tf.not_equal(label, 0))[:, :2], tf.int32)
  min_val = tf.maximum(tf.constant([0, 0]),
                       tf.reduce_max(object_locations, axis=0) - size)
  offset = tf.concat([min_val, [0]], axis=0)

  return random_crop_tensors(tensors, size, offset)[0]
//  [python]
def serve(self, images):
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]
    Returns:
      Tensor holding classification output logits.
    """
    # Skip image preprocessing when input_type is tflite so it is compatible
    # with TFLite quantization.
    image_info = None
    if self._input_type != 'tflite':
      with tf.device('cpu:0'):
        images_spec = tf.TensorSpec(
            shape=self._input_image_size + [3], dtype=tf.float32)
        image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)

        images, image_info = tf.nest.map_structure(
            tf.identity,
            tf.map_fn(
                self._build_inputs,
                elems=images,
                fn_output_signature=(images_spec, image_info_spec),
                parallel_iterations=32))

    outputs = self.inference_step(images)

    # Optionally resize prediction to the input image size.
    if self.params.task.export_config.rescale_output:
      logits = outputs['logits']
      if logits.shape[0] != 1:
        raise ValueError('Batch size cannot be more than 1.')

      image_shape = tf.cast(image_info[0, 0, :], tf.int32)
      if self.params.task.train_data.preserve_aspect_ratio:
        rescale_size = tf.cast(
            tf.math.ceil(image_info[0, 1, :] / image_info[0, 2, :]), tf.int32)
        offsets = tf.cast(image_info[0, 3, :], tf.int32)
        logits = tf.image.resize(logits, rescale_size, method='bilinear')
        outputs['logits'] = tf.image.crop_to_bounding_box(
            logits, offsets[0], offsets[1], image_shape[0], image_shape[1])
      else:
        outputs['logits'] = tf.image.resize(
            logits, [image_shape[0], image_shape[1]], method='bilinear')
    else:
      outputs['logits'] = tf.image.resize(
          outputs['logits'], self._input_image_size, method='bilinear')

    if image_info is not None:
      outputs.update({'image_info': image_info})

    return outputs
//  [python]
def get_resize_random_scale_with_min_size_factor(img_h, img_w, min_size, min_scale=0.7, max_scale=1.3):
  assert min_size is not None
  shorter_side = tf.minimum(img_h, img_w)
  min_scale_factor = t.cast(min_size, t.float32) / t.cast(shorter_side, t.float32)
  min_scale = tf.maximum(min_scale, min_scale_factor)
  max_scale = tf.maximum(max_scale, min_scale_factor)
  scale_factor = tf.random_uniform(shape=[], minval=min_scale, maxval=max_scale, dtype=tf.float32)
  scaled_size = t.cast(tf.round(t.cast([img_h, img_w], t.float32) * scale_factor), t.int32)
  return scaled_size, scale_factor
//  [python]
def get_resize_random_scale_with_min_size_factor(img_h, img_w, min_size, min_scale=0.7, max_scale=1.3):
  assert min_size is not None
  shorter_side = tf.minimum(img_h, img_w)
  min_scale_factor = tf.cast(min_size, tf.float32) / tf.cast(shorter_side, tf.float32)
  min_scale = tf.maximum(min_scale, min_scale_factor)
  max_scale = tf.maximum(max_scale, min_scale_factor)
  scale_factor = tf.random_uniform(shape=[], minval=min_scale, maxval=max_scale, dtype=tf.float32)
  scaled_size = tf.cast(tf.round(tf.cast([img_h, img_w], tf.float32) * scale_factor), tf.int32)
  return scaled_size, scale_factor
//  [python]
def scale_with_min_size(tensors, min_size, min_scale=0.7, max_scale=1.3):
  assert min_size is not None
  img = tensors[DataKeys.IMAGES]

  h = tf.shape(img)[0]
  w = tf.shape(img)[1]
  shorter_side = tf.minimum(h, w)
  min_scale_factor = t.float32(min_size) / t.float32(shorter_side)
  scaled_size = t.int32(t.round(t.float32(tf.shape(img)[:2]) * min_scale_factor))
  tensors_out = resize_fixed_size(tensors, scaled_size)
  return tensors_out
//  [python]
def scale_with_min_size(tensors, min_size, min_scale=0.7, max_scale=1.3):
  assert min_size is not None
  img = tensors[DataKeys.IMAGES]

  h = tf.shape(img)[0]
  w = tf.shape(img)[1]
  shorter_side = tf.minimum(h, w)
  min_scale_factor = tf.cast(min_size, tf.float32) / tf.cast(shorter_side, tf.float32)
  scaled_size = tf.cast(tf.round(tf.cast(tf.shape(img)[:2], tf.float32) * min_scale_factor), tf.int32)
  tensors_out = resize_fixed_size(tensors, scaled_size)
  return tensors_out
//  [python]
def resize_min_short_edge_max_long_edge(tensors, size):
  # Reference: CustomResize._get_augment_params in
  # https://github.com/ppwwyyxx/tensorpack/blob/master/examples/FasterRCNN/common.py
  assert len(size) == 2
  min_short_edge = float(size[0])
  max_long_edge = float(size[1])
  h = t.cast(t.shape(tensors[DataKeys.IMAGES])[0], t.float32)
  w = t.cast(t.shape(tensors[DataKeys.IMAGES])[1], t.float32)

  scale = min_short_edge * 1.0 / t.minimum(h, w)
  less_val = t.less(h, w)
  newh = t.where(less_val, min_short_edge, t.multiply(scale, h))
  neww = t.where(less_val, t.multiply(scale, w), min_short_edge)

  scale = max_long_edge * 1.0 / t.maximum(newh, neww)
  greater_val = t.greater(t.maximum(newh, neww), max_long_edge)
  newh = t.where(greater_val, t.multiply(scale, newh), newh)
  neww = t.where(greater_val, t.multiply(scale, neww), neww)

  newh = t.cast(t.round(newh), t.int32)
  neww = t.cast(t.round(neww), t.int32)

  return resize_fixed_size(tensors, t.stack([newh, neww], axis=0))
//  [python]
def resize_min_short_edge_max_long_edge(tensors, size):
  # Reference: CustomResize._get_augment_params in
  # https://github.com/ppwwyyxx/tensorpack/blob/master/examples/FasterRCNN/common.py
  assert len(size) == 2
  min_short_edge = float(size[0])
  max_long_edge = float(size[1])
  h = tf.cast(tf.shape(tensors[DataKeys.IMAGES])[0], tf.float32)
  w = tf.cast(tf.shape(tensors[DataKeys.IMAGES])[1], tf.float32)

  scale = min_short_edge * 1.0 / tf.minimum(h, w)
  less_val = tf.less(h, w)
  newh = tf.where(less_val, min_short_edge, tf.multiply(scale, h))
  neww = tf.where(less_val, tf.multiply(scale, w), min_short_edge)

  scale = max_long_edge * 1.0 / tf.maximum(newh, neww)
  greater_val = tf.greater(tf.maximum(newh, neww), max_long_edge)
  newh = tf.where(greater_val, tf.multiply(scale, newh), newh)
  neww = tf.where(greater_val, tf.multiply(scale, neww), neww)

  newh = tf.cast(tf.round(newh), tf.int32)
  neww = tf.cast(tf.round(neww), tf.int32)

  return resize_fixed_size(tensors, tf.stack([newh, neww], axis=0))
//  [python]
def test_binary_numerics_with_static(self, op_name):
    np.random.seed(42)
    inputs = (
        np.random.normal(size=(10, 5)).astype(np.float32),
        np.random.normal(size=(10, 5)).astype(np.float32),
    )

    def tf_func(x, y):
      return getattr(tf.raw_ops, op_name)(x=x, y=y)

    if jax.default_backend().lower() == "tpu" and op_name in ("Pow",):
      tols = dict(atol=1e-4)
    else:
      tols = {}

    self._test_convert(tf_func, inputs, **tols)

    # Check static inputs result in static outputs.
    def tf_static():
      vals = tf_func(x=np.array([4.0, 3.0]), y=np.array([1.0, 2.0]))
      return tf.zeros(tf.to_int32(vals))
    self._test_convert(tf_static, [])
//  [python]
def test_binary_numerics_with_static(self, op_name):
    np.random.seed(42)
    inputs = (
        np.random.normal(size=(10, 5)).astype(np.float32),
        np.random.normal(size=(10, 5)).astype(np.float32),
    )

    def tf_func(x, y):
      return getattr(tf.raw_ops, op_name)(x=x, y=y)

    if jax.default_backend().lower() == "tpu" and op_name in ("Pow",):
      tols = dict(atol=1e-4)
    else:
      tols = {}

    self._test_convert(tf_func, inputs, **tols)

    # Check static inputs result in static outputs.
    def tf_static():
      vals = tf_func(x=np.array([4.0, 3.0]), y=np.array([1.0, 2.0]))
      return tf.zeros(tf.cast(vals, tf.int32))
    self._test_convert(tf_static, [])
//  [python]
def test_binary_comparison(self, op_name):
    np.random.seed(42)
    inputs = (np.random.randint(10, size=(10, 5)),
              np.random.randint(10, size=(10, 5)))

    def tf_func(x, y):
      return t.cast(getattr(tf.raw_ops, op_name)(x=x, y=y), tf.int32)
    self._test_convert(tf_func, inputs)
//  [python]
def test_binary_comparison(self, op_name):
    np.random.seed(42)
    inputs = (np.random.randint(10, size=(10, 5)),
              np.random.randint(10, size=(10, 5)))

    def tf_func(x, y):
      return tf.cast(getattr(tf.raw_ops, op_name)(x=x, y=y), tf.int32)
    self._test_convert(tf_func, inputs)
//  [python]
def test_add_n(self):
    np.random.seed(42)
    inputs = [
        [np.random.normal(size=(10, 5)).astype(np.float32) for _ in range(4)]
    ]

    def tf_func(x):
      return tf.raw_ops.AddN(inputs=x)
    self._test_convert(tf_func, inputs)

    # Check static inputs result in static outputs.
    def tf_static():
      vals = tf_func(
          [np.array([4.0, 3.0]), np.array([1.0, 2.0]), np.array([2.0, 1.0])])
      return tf.zeros(tf.to_int32(vals))
    self._test_convert(tf_static, [])
//  [python]
def test_add_n(self):
    np.random.seed(42)
    inputs = [
        [np.random.normal(size=(10, 5)).astype(np.float32) for _ in range(4)]
    ]

    def tf_func(x):
      return tf.raw_ops.AddN(inputs=x)
    self._test_convert(tf_func, inputs)

    # Check static inputs result in static outputs.
    def tf_static():
      vals = tf_func(
          [np.array([4.0, 3.0]), np.array([1.0, 2.0]), np.array([2.0, 1.0])])
      return tf.zeros(tf.cast(vals, tf.int32))
    self._test_convert(tf_static, [])
//  [python]
def test_jitting_static(self, op_name, inputs):
    def tf_func():
      if len(inputs) == 1:
        shape = getattr(tf.raw_ops, op_name)(x=inputs[0])
      else:
        shape = getattr(tf.raw_ops, op_name)(x=inputs[0], y=inputs[1])
      return tf.zeros(tf.to_int32(shape) + 1)

    self._test_convert(tf_func, [])
//  [python]
def test_jitting_static(self, op_name, inputs):
    def tf_func():
      if len(inputs) == 1:
        shape = getattr(tf.raw_ops, op_name)(x=inputs[0])
      else:
        shape = getattr(tf.raw_ops, op_name)(x=inputs[0], y=inputs[1])
      return tf.zeros(tf.cast(shape, tf.int32) + 1)

    self._test_convert(tf_func, [])
//  [python]
def observation(self, t):
    return tf.constant(
        [[1.0, 2.0, 3.0], [0.0, 4.0, 5.0]], dtype=tf.float32
    ) + tf.reshape(t, [1, 1])
//  [python]
def observation(self, t):
    return tf.constant(
        [[1.0, 2.0, 3.0], [0.0, 4.0, 5.0]], dtype=tf.float32
    ) + tf.reshape(tf.cast(t, dtype=tf.float32), [1, 1])
//  [python]
def reward(self, observation, t):
    return tf.concat([observation, tf.zeros([2, 2])], axis=1) - tf.reshape(
        tf.cast(t, dtype=tf.float32), [1, 1]
    )
//  [python]
def _apply_func_with_prob(func: Any,
                          image: tf.Tensor,
                          args: Any,
                          prob: float):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)

  # Apply the function with probability `prob`.
  should_apply_op = tf.floor(tf.random.uniform([], dtype=tf.float32) + prob) > 0
  augmented_image = tf.cond(
      should_apply_op,
      lambda: func(image, *args),
      lambda: image)
  return augmented_image