jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_fake_data(self, inputs, mlm_logits):
        """Sample from the generator to create corrupted input."""
        inputs = pretrain_utils.unmask(inputs)
        disallow = tf.one_hot(
            inputs.masked_lm_ids, depth=self.disc_config.vocab_size,
            dtype=tf.float32) if self._config.disallow_correct else None
        sampled_tokens = tf.stop_gradient(pretrain_utils.sample_from_softmax(
            mlm_logits / self._config.temperature, disallow=disallow))
        sampled_tokids = tf.argmax(sampled_tokens, -1, output_type=tf.int32)
        updated_input_ids, masked = pretrain_utils.scatter_update(
            inputs.input_ids, sampled_tokids, inputs.masked_lm_positions)
        labels = masked * (1 - tf.cast(
            tf.equal(updated_input_ids, inputs.input_ids), tf.int32))
        updated_inputs = pretrain_utils.get_updated_inputs(
            inputs, input_ids=updated_input_ids)
        FakedData = collections.namedtuple("FakedData", [
            "inputs", "is_fake_tokens", "sampled_tokens"])
        return FakedData(inputs=updated_inputs, is_fake_tokens=labels,
                         sampled_tokens=sampled_tokens)
//  [python]
def macro_f1(y, y_hat, thresh=0.5):
    """Compute the macro F1-score on a batch of observations (average F1 across labels)
    Args:
        y: (int32 Tensor)
            labels array of shape (BATCH_SIZE, N_LABELS)
        y_hat: (float32 Tensor)
            probability matrix from forward propagation of shape (BATCH_SIZE, N_LABELS)
        thresh: (float)
            probability value above which we predict positive

    Return
    ------
        macro_f1: (scalar Tensor)
            value of macro F1 for the batch
    """
    y_pred = tf.to_float(tf.greater(y_hat, thresh))
    tp = tf.to_float(tf.math.count_nonzero(y_pred * y, axis=0))
    fp = tf.to_float(tf.math.count_nonzero(y_pred * (1 - y), axis=0))
    fn = tf.to_float(tf.math.count_nonzero((1 - y_pred) * y, axis=0))
    f1 = 2*tp / (2*tp + fn + fp + 1e-16)
    macro_f1 = tf.reduce_mean(f1, axis=-1)
    return macro_f1
//  [python]
def macro_f1(y, y_hat, thresh=0.5):
    """Compute the macro F1-score on a batch of observations (average F1 across labels)
    Args:
        y: (int32 Tensor)
            labels array of shape (BATCH_SIZE, N_LABELS)
        y_hat: (float32 Tensor)
            probability matrix from forward propagation of shape (BATCH_SIZE, N_LABELS)
        thresh: (float)
            probability value above which we predict positive

    Return
    ------
        macro_f1: (scalar Tensor)
            value of macro F1 for the batch
    """
    y_pred = tf.cast(tf.greater(y_hat, thresh), tf.float32)
    tp = tf.cast(tf.math.count_nonzero(y_pred * y, axis=0), tf.float32)
    fp = tf.cast(tf.math.count_nonzero(y_pred * (1 - y), axis=0), tf.float32)
    fn = tf.cast(tf.math.count_nonzero((1 - y_pred) * y, axis=0), tf.float32)
    f1 = 2*tp / (2*tp + fn + fp + 1e-16)
    macro_f1 = tf.reduce_mean(f1, axis=-1)
    return macro_f1
//  [python]
def gaussian_mixture_log_prob(points,
                              params,
                              eps):
  """Computes the likelihood of `points` given GMM `params`."""
  tf.debugging.assert_shapes([
      (points, ('batch_size', 2)),
      (params, ('batch_size', 'n_comp', 4)),  # [comp_weight, mu_x, mu_y, cov]
      (eps, ())
  ])
  mix_probs = params[:, :, 0]  # [b, c]
  mu_s = params[:, :, 1:3]     # [b, c, 2]
  covs = tf.sqrt(params[:, :, 3])  # [b, c]
  diag_s = tf.repeat(covs[:, :, None], repeats=2, axis=2)  # [b, c, 2]
  norm = tfp.distributions.MultivariateNormalDiag(loc=mu_s, scale_diag=diag_s)
  test_points = tf.repeat(points[:, None, :], params.shape[1], 1)
  gauss_log_probs = norm.log_prob(t.cast(test_points, tf.float64))    # [b, c]
  log_gauss_plus_mix = gauss_log_probs + tf.math.log(mix_probs + eps)  # [b, c]
  final_log_probs = tf.reduce_logsumexp(log_gauss_plus_mix, axis=1)    # [b, ]
  return t.cast(final_log_probs, tf.float32)
//  [python]
def gaussian_mixture_log_prob(points,
                              params,
                              eps):
  """Computes the likelihood of `points` given GMM `params`."""
  tf.debugging.assert_shapes([
      (points, ('batch_size', 2)),
      (params, ('batch_size', 'n_comp', 4)),  # [comp_weight, mu_x, mu_y, cov]
      (eps, ())
  ])
  mix_probs = params[:, :, 0]  # [b, c]
  mu_s = params[:, :, 1:3]     # [b, c, 2]
  covs = tf.sqrt(params[:, :, 3])  # [b, c]
  diag_s = tf.repeat(covs[:, :, None], repeats=2, axis=2)  # [b, c, 2]
  norm = tfp.distributions.MultivariateNormalDiag(loc=mu_s, scale_diag=diag_s)
  test_points = tf.repeat(points[:, None, :], params.shape[1], 1)
  gauss_log_probs = norm.log_prob(tf.cast(test_points, tf.float64))    # [b, c]
  log_gauss_plus_mix = gauss_log_probs + tf.math.log(mix_probs + eps)  # [b, c]
  final_log_probs = tf.reduce_logsumexp(log_gauss_plus_mix, axis=1)    # [b, ]
  return tf.cast(final_log_probs, tf.float32)
//  [python]
def project_3d_tf(
    points,
    cam_pose,
    cam_rot,
    ffpp,
    kd,
    weak = False,
    eps = 1e-8):
  """Apply full-perspective projection. Use mean depth if weak=True."""
  points_cent = (points - cam_pose[None])
  cam_rot_mat_t = t.transpose(cam_rot)
  cam_frame_points = points_cent @ cam_rot_mat_t
  xy_3d, zz = cam_frame_points[:, 0:2], cam_frame_points[:, 2, None]
  zz_div = t.reduce_mean(zz, keepdims=True) if weak else zz

  k_mat = t.cast(t.convert_to_tensor([ffpp[:2]]), points.dtype)
  pp = t.cast(t.convert_to_tensor([ffpp[2:]]), points.dtype)

  # order of application as in:
  # http://www.vision.caltech.edu/bouguetj/calib_doc/htmls/parameters.html
  xy = xy_3d / (zz_div + eps)
  xy = apply_distortion(xy, kd)
  xy = xy * k_mat
  xy = xy + pp

  return xy, zz[:, 0]
//  [python]
def project_3d_tf(
    points,
    cam_pose,
    cam_rot,
    ffpp,
    kd,
    weak = False,
    eps = 1e-8):
  """Apply full-perspective projection. Use mean depth if weak=True."""
  points_cent = (points - cam_pose[None])
  cam_rot_mat_t = tf.transpose(cam_rot)
  cam_frame_points = points_cent @ cam_rot_mat_t
  xy_3d, zz = cam_frame_points[:, 0:2], cam_frame_points[:, 2, None]
  zz_div = tf.reduce_mean(zz, keepdims=True) if weak else zz

  k_mat = tf.cast(tf.convert_to_tensor([ffpp[:2]]), points.dtype)
  pp = tf.cast(tf.convert_to_tensor([ffpp[2:]]), points.dtype)

  # order of application as in:
  # http://www.vision.caltech.edu/bouguetj/calib_doc/htmls/parameters.html
  xy = xy_3d / (zz_div + eps)
  xy = apply_distortion(xy, kd)
  xy = xy * k_mat
  xy = xy + pp

  return xy, zz[:, 0]
//  [python]
def get_single_fake_perfect_epi_init(
    data_rec, cam_id):

  """Estimate what a perfect monocular 3D prediction would have been."""
  cams = [data_rec[k][cam_id]
          for k in ['cam_pose3d', 'cam_rot', 'cam_intr', 'cam_kd']]
  bbox = data_rec['bboxes'][cam_id]

  f_mean = cams[2][:2].numpy().mean()
  proj_weak, zz = project_3d_tf(data_rec['pose3d'], *cams, weak=True)
  mean_z = zz.numpy().mean()
  proj_weak = proj_weak.numpy()
  z_scaled = zz[:, None].numpy() / mean_z * f_mean
  cam_frame_xyz = t.concat([proj_weak, z_scaled], axis=1)
  cam_frame_xyz = t.cast(cam_frame_xyz, t.float32)

  size = t.math.maximum(bbox[1] - bbox[0], bbox[3] - bbox[2])
  origin = t.stack([bbox[2], bbox[0], 0], axis=-1)
  size, origin = [t.cast(x, t.float32) for x in [size, origin]]
  unit_xyz = (cam_frame_xyz - origin) / size
  return unit_xyz
//  [python]
def get_single_fake_perfect_epi_init(
    data_rec, cam_id):

  """Estimate what a perfect monocular 3D prediction would have been."""
  cams = [data_rec[k][cam_id]
          for k in ['cam_pose3d', 'cam_rot', 'cam_intr', 'cam_kd']]
  bbox = data_rec['bboxes'][cam_id]

  f_mean = cams[2][:2].numpy().mean()
  proj_weak, zz = project_3d_tf(data_rec['pose3d'], *cams, weak=True)
  mean_z = zz.numpy().mean()
  proj_weak = proj_weak.numpy()
  z_scaled = zz[:, None].numpy() / mean_z * f_mean
  cam_frame_xyz = tf.concat([proj_weak, z_scaled], axis=1)
  cam_frame_xyz = tf.cast(cam_frame_xyz, tf.float32)

  size = tf.math.maximum(bbox[1] - bbox[0], bbox[3] - bbox[2])
  origin = tf.stack([bbox[2], bbox[0], 0], axis=-1)
  size, origin = [tf.cast(x, tf.float32) for x in [size, origin]]
  unit_xyz = (cam_frame_xyz - origin) / size
  return unit_xyz
//  [python]
def gather_npu(params, indices):
  def grad(dy):
    params_shape = tf.shape(params, out_type=tf.int64)
    params_shape = t.cast(params_shape, tf.int32)
    grad_gather = tf.unsorted_segment_sum(dy, indices, params_shape[0])
    return grad_gather, None
  return tf.gather(params, indices), grad
//  [python]
def gather_npu(params, indices):
  def grad(dy):
    params_shape = tf.shape(params, out_type=tf.int64)
    params_shape = tf.cast(params_shape, tf.int32)
    grad_gather = tf.unsorted_segment_sum(dy, indices, params_shape[0])
    return grad_gather, None
  return tf.gather(params, indices), grad
//  [python]
def convert_rec_pose2d_to_bbox_axis(
    input_rec):
  """Converts coordinates / mixture params in a record from pixels to [0, 1]."""
  # full keys list:
  # 'pose3d', 'cam_pose3d', 'cam_rot', 'cam_intr', 'cam_kd', 'pose2d_gt',
  # 'pose2d_repr', 'heatmaps', 'pose2d_pred', 'keys',
  # 'bboxes', 'pose3d_epi_pred'

  bboxes, pose2d_gt, pose2d_repr, mix_params, mean_pred2d = [
      input_rec[x] for x in
      ['bboxes', 'pose2d_gt', 'pose2d_repr', 'heatmaps', 'pose2d_pred']]

  sizes = tf.math.maximum(
      bboxes[:, 1] - bboxes[:, 0], bboxes[:, 3] - bboxes[:, 2])
  origins = tf.stack([bboxes[:, 2], bboxes[:, 0]], axis=-1)

  sizes, origins = [tf.to_float(x) for x in [sizes, origins]]

  shifted_mixture_means = mix_params[:, :, :, 1:3] - origins[:, None, None, :]
  mix_params_new = tf.concat([
      mix_params[:, :, :, 0, None],
      (shifted_mixture_means / sizes[:, None, None, None]),
      mix_params[:, :, :, 3, None] / sizes[:, None, None, None]**2
  ], axis=-1)

  pose2d_gt_new = (pose2d_gt - origins[:, None, :]) / sizes[:, None, None]
  pose2d_repr_new = (pose2d_repr - origins[:, None, :]) / sizes[:, None, None]
  mean_pred_new = (mean_pred2d - origins[:, None, :]) / sizes[:, None, None]

  override_rec = {
      'pose2d_gt': pose2d_gt_new,
      'pose2d_repr': pose2d_repr_new,
      'heatmaps': mix_params_new,
      'pose2d_pred': mean_pred_new,
  }

  assert not(override_rec.keys() - input_rec.keys()), 'override has new keys'
  return {**input_rec, **override_rec}
//  [python]
def convert_rec_pose2d_to_bbox_axis(
    input_rec):
  """Converts coordinates / mixture params in a record from pixels to [0, 1]."""
  # full keys list:
  # 'pose3d', 'cam_pose3d', 'cam_rot', 'cam_intr', 'cam_kd', 'pose2d_gt',
  # 'pose2d_repr', 'heatmaps', 'pose2d_pred', 'keys',
  # 'bboxes', 'pose3d_epi_pred'

  bboxes, pose2d_gt, pose2d_repr, mix_params, mean_pred2d = [
      input_rec[x] for x in
      ['bboxes', 'pose2d_gt', 'pose2d_repr', 'heatmaps', 'pose2d_pred']]

  sizes = tf.math.maximum(
      bboxes[:, 1] - bboxes[:, 0], bboxes[:, 3] - bboxes[:, 2])
  origins = tf.stack([bboxes[:, 2], bboxes[:, 0]], axis=-1)

  sizes, origins = [tf.cast(x, tf.float64) for x in [sizes, origins]]

  shifted_mixture_means = mix_params[:, :, :, 1:3] - origins[:, None, None, :]
  mix_params_new = tf.concat([
      mix_params[:, :, :, 0, None],
      (shifted_mixture_means / sizes[:, None, None, None]),
      mix_params[:, :, :, 3, None] / sizes[:, None, None, None]**2
  ], axis=-1)

  pose2d_gt_new = (pose2d_gt - origins[:, None, :]) / sizes[:, None, None]
  pose2d_repr_new = (pose2d_repr - origins[:, None, :]) / sizes[:, None, None]
  mean_pred_new = (mean_pred2d - origins[:, None, :]) / sizes[:, None, None]

  override_rec = {
      'pose2d_gt': pose2d_gt_new,
      'pose2d_repr': pose2d_repr_new,
      'heatmaps': mix_params_new,
      'pose2d_pred': mean_pred_new,
  }

  assert not(override_rec.keys() - input_rec.keys()), 'override has new keys'
  return {**input_rec, **override_rec}
//  [python]
def embedding_lookup_npu(params, indices):
  def grad(dy):
    params_shape = tf.shape(params, out_type=tf.int64)
    params_shape = tf.to_int32(params_shape)
    grad_embedding_lookup = tf.unsorted_segment_sum(dy, indices, params_shape[0])
    return grad_embedding_lookup, None
  return tf.nn.embedding_lookup(params, indices), grad
//  [python]
def embedding_lookup_npu(params, indices):
  def grad(dy):
    params_shape = tf.shape(params, out_type=tf.int64)
    params_shape = tf.cast(params_shape, tf.int32)
    grad_embedding_lookup = tf.unsorted_segment_sum(dy, indices, params_shape[0])
    return grad_embedding_lookup, None
  return tf.nn.embedding_lookup(params, indices), grad
//  [python]
def createModelWithMultipleMixedInputs(self, save_as_keras):
    dense_input = tf_keras.layers.Input(
        shape=(2,), name='input_1', dtype=tf.int64
    )
    dense_float_input = tf.cast(dense_input, tf.float32)
    sparse_input = tf_keras.layers.Input(
        shape=(1,), name='input_2', sparse=True
    )
    dense_sparse_input = tf_keras.layers.Dense(1, name='dense_input2')(
        sparse_input
    )
    ragged_input = tf_keras.layers.Input(
        shape=(None,), name='input_3', ragged=True
    )
    dense_ragged_input = tf_keras.layers.Lambda(lambda x: x.to_tensor())(
        ragged_input
    )
    dense_ragged_input.set_shape((None, 1))
    inputs = [dense_input, sparse_input, ragged_input]
    input_layer = tf_keras.layers.concatenate(
        [dense_float_input, dense_sparse_input, dense_ragged_input]
    )
    output_layer = tf_keras.layers.Dense(1, activation=tf.nn.sigmoid)(
        input_layer
    )
    model = tf_keras.models.Model(inputs, output_layer)

    @tf.function
    def serving_default(features):
      return model(features)

    input_spec = {
        'input_1':
            tf.TensorSpec(shape=(None, 2), dtype=tf.int64, name='input_1'),
        'input_2':
            tf.SparseTensorSpec(shape=(None, 1), dtype=tf.float32),
        'input_3':
            tf.RaggedTensorSpec(shape=(None, 1), dtype=tf.float32)
    }
    signatures = {
        'serving_default': serving_default.get_concrete_function(input_spec),
        'custom_signature': serving_default.get_concrete_function(input_spec),
    }

    export_path = tempfile.mkdtemp()
    if save_as_keras:
      model.save(export_path, save_format='tf', signatures=signatures)
    else:
      tf.saved_model.save(model, export_path, signatures=signatures)
    return export_path
//  [python]
def auxiliary_loss(self, h_states, click_seq, noclick_seq, mask, stag = None):
        mask = tf.cast(mask, tf.float32)
        click_input_ = tf.concat([h_states, click_seq], -1)
        noclick_input_ = tf.concat([h_states, noclick_seq], -1)
        click_prop_ = self.auxiliary_net(click_input_, stag = stag)[:, :, 0]
        noclick_prop_ = self.auxiliary_net(noclick_input_, stag = stag)[:, :, 0]
        click_loss_ = - tf.reshape(tf.log(click_prop_), [-1, tf.shape(click_seq)[1]]) * mask
        noclick_loss_ = - tf.reshape(tf.log(1.0 - noclick_prop_), [-1, tf.shape(noclick_seq)[1]]) * mask
        loss_ = tf.reduce_sum(click_loss_ + noclick_loss_) / (tf.reduce_sum(mask) + 1e-7)
        return loss_
//  [python]
def compute_opt_stats(input_rec,
                      iter_opt_results
                      ):
  """Computes per-step metrics for the output of `optimize_heatmap_logp`.

  See `run_inference_optimization` for the full spec of inputs and outputs.

  Args:
    input_rec: an input dict of heatmaps, monocular 3D poses, etc.
    iter_opt_results: an output dict produced by optimize_heatmap_logp.

  Returns:
    A dict of tensors combining input record with predictions and metrics.
  """
  pose3d_gt, pose2d_gt, mean_posenet_pred2d = [
      tf.cast(input_rec[x], tf.float32)
      for x in ['pose3d', 'pose2d_gt', 'pose2d_pred']
  ]

  iters, losses, *opt_results = iter_opt_results
  n_report = losses.shape[0]

  # [n_report, 4, 17, 2]
  opt_pose2d_preds = tf.convert_to_tensor([
      project3d_weak(*[x[viz_id] for x in opt_results])[Ellipsis, :2]
      for viz_id in range(n_report)
  ])

  gt_aligned_projected2d = []
  for viz_id in range(n_report):
    gt_aligned_pose = align_aba(pose3d_gt, opt_results[0][viz_id])[0]
    other_opt_params = [x[viz_id] for x in opt_results[1:]]
    projected2d = project3d_weak(gt_aligned_pose, *other_opt_params)
    gt_aligned_projected2d.append(projected2d[Ellipsis, :2])

  gt_aligned_projected2d = tf.convert_to_tensor(
      gt_aligned_projected2d, dtype=tf.float32)

  iter_pmpjpe = batch_pmpjpe(opt_results[0], pose3d_gt)

  iter_pose2d_err = mean_norm(opt_pose2d_preds - pose2d_gt, 3, (1, 2))
  mean_posenet_gt_err = mean_norm(pose2d_gt - mean_posenet_pred2d, -1, None)
  iter_mean_posenet_err = mean_norm(
      opt_pose2d_preds - mean_posenet_pred2d, 3, (1, 2))
  iter_gt2d_gt_aligned_proj_err = mean_norm(
      gt_aligned_projected2d - pose2d_gt, 3, (1, 2))

  augment_rec = {
      'loss': losses,
      'iters': iters,
      'pose3d_opt_preds': opt_results[0],
      'cam_rot_opt_preds': opt_results[1],
      'scale_opt_preds': opt_results[2],
      'shift_opt_preds': opt_results[3],

      'pose2d_opt_preds': opt_pose2d_preds,
      'pose3d_gt_aligned_pred_3d_proj': gt_aligned_projected2d,
      'pose3d_pred_pmpjpe': iter_pmpjpe,
      'pose2d_pred_err': iter_pose2d_err,
      'pose2d_pred_vs_posenet_err': iter_mean_posenet_err,
      'pose2d_gt_posenet_err_mean': mean_posenet_gt_err,
      'pose3d_gt_backaligned_pose2d_gt_err': iter_gt2d_gt_aligned_proj_err,
  }

  assert not(augment_rec.keys() & input_rec.keys()), 'augment overrides keys'
  return {**augment_rec, **input_rec}
//  [python]
def map_layer_inputs_to_op(cls, z, rec_weights, i, initial_state=None):
        """
        Just like NativeOp.LstmGenericBase.map_layer_inputs_to_op().

        :param tf.Tensor z: Z: inputs: shape (time,batch,n_hidden*4)
        :param tf.Tensor rec_weights: V_h / W_re: shape (n_hidden,n_hidden*4)
        :param tf.Tensor i: index: shape (time,batch)
        :param tf.Tensor|None initial_state: shape (batch,n_hidden)
        :rtype: (tf.Tensor,tf.Tensor,tf.Tensor,tf.Tensor)
        """
        assert z.get_shape().ndims == 3
        assert rec_weights.get_shape().ndims == 2
        assert i.get_shape().ndims == 2
        if i.dtype != tf.float32:
            if not hasattr(i, "cast_float32"):
                from returnn.tf.util.basic import reuse_name_scope_of_tensor

                with reuse_name_scope_of_tensor(i):
                    i_cast_float32 = tf.to_float(i, name="index_cast_float32")
                i.cast_float32 = i_cast_float32
            i = i.cast_float32
        n_batch = tf.shape(z)[1]
        n_out = tf.shape(rec_weights)[0]
        if initial_state is not None:
            if isinstance(initial_state, rnn_cell.LSTMStateTuple):
                initial_state = initial_state.c
            c = initial_state
        else:
            c = tf.zeros((n_batch, n_out), dtype=tf.float32)
        return z, rec_weights, c, i
//  [python]
def map_layer_inputs_to_op(cls, z, rec_weights, i, initial_state=None):
        """
        Just like NativeOp.LstmGenericBase.map_layer_inputs_to_op().

        :param tf.Tensor z: Z: inputs: shape (time,batch,n_hidden*4)
        :param tf.Tensor rec_weights: V_h / W_re: shape (n_hidden,n_hidden*4)
        :param tf.Tensor i: index: shape (time,batch)
        :param tf.Tensor|None initial_state: shape (batch,n_hidden)
        :rtype: (tf.Tensor,tf.Tensor,tf.Tensor,tf.Tensor)
        """
        assert z.get_shape().ndims == 3
        assert rec_weights.get_shape().ndims == 2
        assert i.get_shape().ndims == 2
        if i.dtype != tf.float32:
            if not hasattr(i, "cast_float32"):
                from returnn.tf.util.basic import reuse_name_scope_of_tensor

                with reuse_name_scope_of_tensor(i):
                    i_cast_float32 = tf.cast(i, dtype=tf.float32, name="index_cast_float32")
                i.cast_float32 = i_cast_float32
            i = i.cast_float32
        n_batch = tf.shape(z)[1]
        n_out = tf.shape(rec_weights)[0]
        if initial_state is not None:
            if isinstance(initial_state, rnn_cell.LSTMStateTuple):
                initial_state = initial_state.c
            c = initial_state
        else:
            c = tf.zeros((n_batch, n_out), dtype=tf.float32)
        return z, rec_weights, c, i
//  [python]
def _parse_eval_data(self, data):
    """Parses data for training and evaluation."""
    classes = data['groundtruth_classes'] + self._label_shift
    boxes = data['groundtruth_boxes']
    is_crowd = data['groundtruth_is_crowd']

    # Gets original image and its size.
    image = data['image']
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)

    image_shape = tf.shape(image)[:2]
    boxes = box_ops.denormalize_boxes(boxes, image_shape)
    gt_boxes = boxes
    image, image_info = preprocess_ops.resize_image(
        image, min(self._output_size), max(self._output_size)
    )
    boxes = preprocess_ops.resize_and_crop_boxes(
        boxes, image_info[2, :], image_info[1, :], image_info[3, :]
    )
    scale = tf.concat([self._output_size, self._output_size], -1)
    boxes = boxes / scale

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

    prompt_seq = tf.constant([pix2seq_cfg.OD_ID], dtype=tf.int64)
    backgrnd_val = 0.3
    image = backgrnd_val + tf.image.pad_to_bounding_box(
        image - backgrnd_val, 0, 0, self._output_size[0], self._output_size[1]
    )

    labels = {
        'prompt': prompt_seq,
        'classes': preprocess_ops.clip_or_pad_to_fixed_size(
            classes, self._max_num_boxes
        ),
        'boxes': preprocess_ops.clip_or_pad_to_fixed_size(
            boxes, self._max_num_boxes
        ),
    }
    labels.update({
        'id': int(data['source_id']),
        'image_info': image_info,
        'is_crowd': preprocess_ops.clip_or_pad_to_fixed_size(
            is_crowd, self._max_num_boxes
        ),
        'gt_boxes': preprocess_ops.clip_or_pad_to_fixed_size(
            gt_boxes, self._max_num_boxes
        ),
    })

    return image, labels
//  [python]
def _parse_eval_data(self, data):
    """Parses data for training and evaluation."""
    classes = data['groundtruth_classes'] + self._label_shift
    boxes = data['groundtruth_boxes']
    is_crowd = data['groundtruth_is_crowd']

    # Gets original image and its size.
    image = data['image']
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)

    image_shape = tf.shape(image)[:2]
    boxes = box_ops.denormalize_boxes(boxes, image_shape)
    gt_boxes = boxes
    image, image_info = preprocess_ops.resize_image(
        image, min(self._output_size), max(self._output_size)
    )
    boxes = preprocess_ops.resize_and_crop_boxes(
        boxes, image_info[2, :], image_info[1, :], image_info[3, :]
    )
    scale = tf.cast(
        tf.concat([self._output_size, self._output_size], -1), boxes.dtype
    )
    boxes = boxes / scale

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

    prompt_seq = tf.constant([pix2seq_cfg.OD_ID], dtype=tf.int64)
    backgrnd_val = 0.3
    image = backgrnd_val + tf.image.pad_to_bounding_box(
        image - backgrnd_val, 0, 0, self._output_size[0], self._output_size[1]
    )

    labels = {
        'prompt': prompt_seq,
        'classes': preprocess_ops.clip_or_pad_to_fixed_size(
            classes, self._max_num_boxes
        ),
        'boxes': preprocess_ops.clip_or_pad_to_fixed_size(
            boxes, self._max_num_boxes
        ),
    }
    labels.update({
        'id': int(data['source_id']),
        'image_info': image_info,
        'is_crowd': preprocess_ops.clip_or_pad_to_fixed_size(
            is_crowd, self._max_num_boxes
        ),
        'gt_boxes': preprocess_ops.clip_or_pad_to_fixed_size(
            gt_boxes, self._max_num_boxes
        ),
    })

    return image, labels
//  [python]
def _lstm_last_output(out, y0, mask, start: int, step: Optional[int]):
    if step is None:
        step = 1
    n_time = tf.shape(out)[0]
    if step < 0:
        start = n_time - start - 1
    mask = t.cast(mask, t.int32)  # [T,B]
    mask_indices = tf.range(n_time)[:, None]  # [T,1]
    mask_indices = tf_util.where_bc(t.cast(mask, t.bool), mask_indices, -1 if step > 0 else n_time)  # [T,B]
    mask_indices = mask_indices[start::step]  # [T',B]
    last_indices = (tf.reduce_max if step > 0 else tf.reduce_min)(mask_indices, axis=0)  # [B]
    last_indices_ = tf.clip_by_value(last_indices, clip_value_min=0, clip_value_max=n_time - 1)  # [B]
    idx_ext = tf_util.nd_indices(last_indices_, indices_batch_major=False)  # [B,2] -> (time,batch) indices
    final_output = tf.gather_nd(out, indices=idx_ext)  # [B,n_hidden]
    valid = tf.logical_and(tf.greater_equal(last_indices, 0), tf.less(last_indices, n_time))  # [B]
    final_output = tf_util.where_bc(valid[:, None], final_output, y0)
    return final_output
//  [python]
def _lstm_last_output(out, y0, mask, start: int, step: Optional[int]):
    if step is None:
        step = 1
    n_time = tf.shape(out)[0]
    if step < 0:
        start = n_time - start - 1
    mask = tf.cast(mask, tf.int32)  # [T,B]
    mask_indices = tf.range(n_time)[:, None]  # [T,1]
    mask_indices = tf_util.where_bc(tf.cast(mask, tf.bool), mask_indices, -1 if step > 0 else n_time)  # [T,B]
    mask_indices = mask_indices[start::step]  # [T',B]
    last_indices = (tf.reduce_max if step > 0 else tf.reduce_min)(mask_indices, axis=0)  # [B]
    last_indices_ = tf.clip_by_value(last_indices, clip_value_min=0, clip_value_max=n_time - 1)  # [B]
    idx_ext = tf_util.nd_indices(last_indices_, indices_batch_major=False)  # [B,2] -> (time,batch) indices
    final_output = tf.gather_nd(out, indices=idx_ext)  # [B,n_hidden]
    valid = tf.logical_and(tf.greater_equal(last_indices, 0), tf.less(last_indices, n_time))  # [B]
    final_output = tf_util.where_bc(valid[:, None], final_output, y0)
    return final_output
//  [python]
def update_state(
      self,
      *args: tf.Tensor,
      sample_weight: tf.Tensor | None = None,
      slicing_feature: tf.Tensor,
      **kwargs,
  ):
    """Updates the state of the metrics for each slice.

    Args:
      *args: A variable amount of `tf.Tensor` instances that will be passed to
        the `update_state` method of each metric.
      sample_weight: An optional `tf.Tensor` used to weight the sample. Its
        dimensions must be broadcastable to the shape(s) of *args.
      slicing_feature: A `tf.Tensor` consisting of the feature to be sliced on.
        Its dimensions must be broadcastable to the shape(s) of *args.
      **kwargs: Keyword arguments that will be passed to the `update_state`
        method of each metric.
    """

    if slicing_feature.dtype != self._slicing_feature_dtype:
      raise ValueError(
          "The `slicing_feature` and slicing values in `slicing_spec` must "
          "have the same type. Got types: "
          f"{(slicing_feature.dtype, self._slicing_feature_dtype)}."
      )

    if sample_weight is not None:
      for _ in range(len(slicing_feature.shape) - len(sample_weight.shape)):
        sample_weight = tf.expand_dims(sample_weight, axis=-1)

      for _ in range(len(sample_weight.shape) - len(slicing_feature.shape)):
        slicing_feature = tf.expand_dims(slicing_feature, axis=-1)

    self._metric.update_state(*args, sample_weight=sample_weight, **kwargs)
    for slicing_val, metric in zip(
        self._slicing_values_tensors, self._sliced_metrics
    ):
      slice_mask = tf.cast(slicing_feature == slicing_val, dtype=tf.float32)
      if sample_weight is not None:
        weight = slice_mask * tf.cast(sample_weight, dtype=tf.float32)
      else:
        weight = slice_mask
      metric.update_state(*args, sample_weight=weight, **kwargs)
//  [python]
def map_layer_inputs_to_op(self, x, weights, b, i, initial_state=None):
        """
        Just like NativeOp.LstmGenericBase.map_layer_inputs_to_op().
        :param tf.Tensor x: inputs: shape (time,batch,n_input_dim)
        :param tf.Tensor weights: shape (n_input_dim+n_hidden,n_hidden*4)
        :param tf.Tensor b: shape (n_hidden*4,)
        :param tf.Tensor i: index: shape (time,batch)
        :param tf.Tensor|None initial_state: shape (batch,n_hidden)
        :rtype: tuple[tf.Tensor]
        """
        x.set_shape(tf.TensorShape([None, None, self.n_input_dim]))
        weights.set_shape(tf.TensorShape([self.n_input_dim + self.n_hidden, self.n_hidden * 4]))
        i.set_shape(tf.TensorShape([None, None]))
        if i.dtype != tf.float32:
            if not hasattr(i, "cast_float32"):
                from returnn.tf.util.basic import reuse_name_scope_of_tensor

                with reuse_name_scope_of_tensor(i):
                    i_cast_float32 = tf.to_float(i, name="index_cast_float32")
                i.cast_float32 = i_cast_float32
            i = i.cast_float32
        n_batch = tf.shape(x)[1]
        if initial_state is not None:
            c0 = initial_state
        else:
            c0 = tf.zeros((n_batch, self.n_hidden), dtype=tf.float32, name="initial_c")
        # We could make `h` a variable exactly if `c` is a trainable variable.
        y0 = tf.zeros((n_batch, self.n_hidden), dtype=tf.float32, name="initial_h")
        start = tf.constant(0, name="start")
        step = tf.constant(self.step or 1, name="step")
        return x, weights, b, y0, c0, i, start, step
//  [python]
def map_layer_inputs_to_op(self, x, weights, b, i, initial_state=None):
        """
        Just like NativeOp.LstmGenericBase.map_layer_inputs_to_op().
        :param tf.Tensor x: inputs: shape (time,batch,n_input_dim)
        :param tf.Tensor weights: shape (n_input_dim+n_hidden,n_hidden*4)
        :param tf.Tensor b: shape (n_hidden*4,)
        :param tf.Tensor i: index: shape (time,batch)
        :param tf.Tensor|None initial_state: shape (batch,n_hidden)
        :rtype: tuple[tf.Tensor]
        """
        x.set_shape(tf.TensorShape([None, None, self.n_input_dim]))
        weights.set_shape(tf.TensorShape([self.n_input_dim + self.n_hidden, self.n_hidden * 4]))
        i.set_shape(tf.TensorShape([None, None]))
        if i.dtype != tf.float32:
            if not hasattr(i, "cast_float32"):
                from returnn.tf.util.basic import reuse_name_scope_of_tensor

                with reuse_name_scope_of_tensor(i):
                    i_cast_float32 = tf.cast(i, dtype=tf.float32, name="index_cast_float32")
                i.cast_float32 = i_cast_float32
            i = i.cast_float32
        n_batch = tf.shape(x)[1]
        if initial_state is not None:
            c0 = initial_state
        else:
            c0 = tf.zeros((n_batch, self.n_hidden), dtype=tf.float32, name="initial_c")
        # We could make `h` a variable exactly if `c` is a trainable variable.
        y0 = tf.zeros((n_batch, self.n_hidden), dtype=tf.float32, name="initial_h")
        start = tf.constant(0, name="start")
        step = tf.constant(self.step or 1, name="step")
        return x, weights, b, y0, c0, i, start, step
//  [python]
def chunk(x, index, chunk_size, chunk_step):
    """
    :param tf.Tensor x: (time,batch,dim)
    :param tf.Tensor index: (time,batch)
    :param int|tf.Tensor chunk_size:
    :param int|tf.Tensor chunk_step:
    :return: out, oindex.
      out is of shape (chunk_size, n_batch * n_chunks, n_dim), oindex of shape (chunk_size, n_batch * n_chunks).
    :rtype: (tf.Tensor,tf.Tensor)
    """
    assert x.get_shape().ndims == 3
    x_shape = tf.shape(x)
    n_time = x_shape[0]
    n_batch = x_shape[1]
    n_dim = x_shape[2]
    n_chunks = tf.maximum(n_time - chunk_size + chunk_step - 1, 0) // chunk_step + 1
    chunk_params = t.cast([chunk_size, chunk_step], tf.float32)  # that's currently the API...
    out_buffer = tf.zeros((chunk_size, n_batch * n_chunks, n_dim), dtype=x.dtype)
    oindex_buffer = tf.zeros((chunk_size, n_batch * n_chunks), dtype=index.dtype)
    chunk_op = OpMaker(OpDescription.from_gen_base(native_op.Chunking)).make_op(grad_func=_chunk_grad)
    out, oindex = chunk_op(x, index, out_buffer, oindex_buffer, chunk_params)
    return out, oindex
//  [python]
def chunk(x, index, chunk_size, chunk_step):
    """
    :param tf.Tensor x: (time,batch,dim)
    :param tf.Tensor index: (time,batch)
    :param int|tf.Tensor chunk_size:
    :param int|tf.Tensor chunk_step:
    :return: out, oindex.
      out is of shape (chunk_size, n_batch * n_chunks, n_dim), oindex of shape (chunk_size, n_batch * n_chunks).
    :rtype: (tf.Tensor,tf.Tensor)
    """
    assert x.get_shape().ndims == 3
    x_shape = tf.shape(x)
    n_time = x_shape[0]
    n_batch = x_shape[1]
    n_dim = x_shape[2]
    n_chunks = tf.maximum(n_time - chunk_size + chunk_step - 1, 0) // chunk_step + 1
    chunk_params = tf.cast([chunk_size, chunk_step], tf.float32)  # that's currently the API...
    out_buffer = tf.zeros((chunk_size, n_batch * n_chunks, n_dim), dtype=x.dtype)
    oindex_buffer = tf.zeros((chunk_size, n_batch * n_chunks), dtype=index.dtype)
    chunk_op = OpMaker(OpDescription.from_gen_base(native_op.Chunking)).make_op(grad_func=_chunk_grad)
    out, oindex = chunk_op(x, index, out_buffer, oindex_buffer, chunk_params)
    return out, oindex
//  [python]
def _unchunk_grad(op, *output_grads):
    """
    :param tf.Operation op:
    :param tf.Tensor output_grads:
    :return: input_grads
    :rtype: tuple[tf.Tensor]
    """
    x, index, _, _, _, chunk_params = op.inputs
    out_grad, _, _ = output_grads
    chunk_size = t.cast(chunk_params[0], t.int32)
    chunk_step = t.cast(chunk_params[1], t.int32)
    out, oindex, factors = op.outputs
    out_grad *= tf.expand_dims(factors, axis=2)
    x_grad, _ = chunk(out_grad, index=oindex, chunk_size=chunk_size, chunk_step=chunk_step)
    return x_grad, None, None, None, None, None
//  [python]
def _unchunk_grad(op, *output_grads):
    """
    :param tf.Operation op:
    :param tf.Tensor output_grads:
    :return: input_grads
    :rtype: tuple[tf.Tensor]
    """
    x, index, _, _, _, chunk_params = op.inputs
    out_grad, _, _ = output_grads
    chunk_size = tf.cast(chunk_params[0], tf.int32)
    chunk_step = tf.cast(chunk_params[1], tf.int32)
    out, oindex, factors = op.outputs
    out_grad *= tf.expand_dims(factors, axis=2)
    x_grad, _ = chunk(out_grad, index=oindex, chunk_size=chunk_size, chunk_step=chunk_step)
    return x_grad, None, None, None, None, None
//  [python]
def build(self, D, M, K):
        # params
        self.W = tf.Variable(tf.random.normal(shape=(D, K, M)) * np.sqrt(2.0 / M))
        self.c = tf.Variable(np.zeros(M).astype(np.float32))
        self.b = tf.Variable(np.zeros((D, K)).astype(np.float32))

        # data
        self.X_in = tf.compat.v1.placeholder(tf.float32, shape=(None, D))

        # one hot encode X
        # first, make each rating an int
        X = tf.cast(self.X_in * 2 - 1, tf.int32)
        X = tf.one_hot(X, K)

        # conditional probabilities
        # NOTE: tf.contrib.distributions.Bernoulli API has changed in Tensorflow v1.2
        V = X
        p_h_given_v = tf.nn.sigmoid(dot1(V, self.W) + self.c)
        self.p_h_given_v = p_h_given_v # save for later

        # draw a sample from p(h | v)
        r = tf.random.uniform(shape=tf.shape(input=p_h_given_v))
        H = tf.cast(r < p_h_given_v, dtype=tf.float32)

        # draw a sample from p(v | h)
        # note: we don't have to actually do the softmax
        logits = dot2(H, self.W) + self.b
        cdist = tf.compat.v1.distributions.Categorical(logits=logits)
        X_sample = cdist.sample() # shape is (N, D)
        X_sample = tf.one_hot(X_sample, depth=K) # turn it into (N, D, K)

        # mask X_sample to remove missing ratings
        mask2d = tf.cast(self.X_in > 0, tf.float32)
        mask3d = tf.stack([mask2d]*K, axis=-1) # repeat K times in last dimension
        X_sample = X_sample * mask3d


        # build the objective
        objective = tf.reduce_mean(input_tensor=self.free_energy(X)) - tf.reduce_mean(input_tensor=self.free_energy(X_sample))
        self.train_op = tf.compat.v1.train.AdamOptimizer(1e-2).minimize(objective)
        # self.train_op = tf.train.GradientDescentOptimizer(1e-3).minimize(objective)

        # build the cost
        # we won't use this to optimize the model parameters
        # just to observe what happens during training
        logits = self.forward_logits(X)
        self.cost = tf.reduce_mean(
            input_tensor=tf.nn.softmax_cross_entropy_with_logits(
                labels=tf.stop_gradient(X),
                logits=logits,
            )
        )

        # to get the output
        self.output_visible = self.forward_output(X)


        # for calculating SSE
        self.one_to_ten = tf.constant((np.arange(10) + 1).astype(np.float32) / 2)
        self.pred = tf.tensordot(self.output_visible, self.one_to_ten, axes=[[2], [0]])
        mask = tf.cast(self.X_in > 0, tf.float32)
        se = mask * (self.X_in - self.pred) * (self.X_in - self.pred)
        self.sse = tf.reduce_sum(input_tensor=se)

        # test SSE
        self.X_test = tf.compat.v1.placeholder(tf.float32, shape=(None, D))
        mask = tf.cast(self.X_test > 0, tf.float32)
        tse = mask * (self.X_test - self.pred) * (self.X_test - self.pred)
        self.tsse = tf.reduce_sum(input_tensor=tse)


        initop = tf.compat.v1.global_variables_initializer()
        self.session = tf.compat.v1.Session()
        self.session.run(initop)
//  [python]
def _chunk_grad(op, *output_grads):
    """
    :param tf.Operation op:
    :param tf.Tensor output_grads:
    :return: input_grads
    :rtype: tuple[tf.Tensor]
    """
    x, index, _, _, chunk_params = op.inputs  # type: tf.Tensor
    out_grad, _ = output_grads
    assert x.get_shape().ndims == 3
    x_shape = tf.shape(x)
    n_time = x_shape[0]
    n_batch = x_shape[1]
    chunk_size = t.cast(chunk_params[0], t.int32)
    chunk_step = t.cast(chunk_params[1], t.int32)
    out, oindex = op.outputs
    x_grad, _, factors = unchunk(
        out_grad, index=oindex, chunk_size=chunk_size, chunk_step=chunk_step, n_time=n_time, n_batch=n_batch
    )
    # We applied the factor in unchunk, but for this gradient, we actually don't want that, so undo it.
    x_grad /= tf.expand_dims(factors, axis=2)
    return x_grad, None, None, None, None
//  [python]
def _chunk_grad(op, *output_grads):
    """
    :param tf.Operation op:
    :param tf.Tensor output_grads:
    :return: input_grads
    :rtype: tuple[tf.Tensor]
    """
    x, index, _, _, chunk_params = op.inputs  # type: tf.Tensor
    out_grad, _ = output_grads
    assert x.get_shape().ndims == 3
    x_shape = tf.shape(x)
    n_time = x_shape[0]
    n_batch = x_shape[1]
    chunk_size = tf.cast(chunk_params[0], tf.int32)
    chunk_step = tf.cast(chunk_params[1], tf.int32)
    out, oindex = op.outputs
    x_grad, _, factors = unchunk(
        out_grad, index=oindex, chunk_size=chunk_size, chunk_step=chunk_step, n_time=n_time, n_batch=n_batch
    )
    # We applied the factor in unchunk, but for this gradient, we actually don't want that, so undo it.
    x_grad /= tf.expand_dims(factors, axis=2)
    return x_grad, None, None, None, None
//  [python]
def unchunk(x, index, chunk_size, chunk_step, n_time, n_batch):
    """
    :param tf.Tensor x: output e.g. from :func:`chunk`
    :param tf.Tensor index:
    :param int|tf.Tensor chunk_size:
    :param int|tf.Tensor chunk_step:
    :param tf.Tensor n_time:
    :param tf.Tensor n_batch:
    :return: out, oindex, ofactors
    :rtype: (tf.Tensor,tf.Tensor,tf.Tensor)
    """
    assert x.get_shape().ndims == 3
    n_dim = tf.shape(x)[2]
    chunk_params = t.cast([chunk_size, chunk_step], t.float32)  # that's currently the API...
    out_buffer = tf.zeros((n_time, n_batch, n_dim), dtype=x.dtype)
    oindex_buffer = tf.zeros((n_time, n_batch), dtype=index.dtype)
    ofactors_buffer = tf.zeros((n_time, n_batch), dtype=x.dtype)
    unchunk_op = OpMaker(OpDescription.from_gen_base(native_op.UnChunking)).make_op(grad_func=_unchunk_grad)
    out, oindex, ofactors = unchunk_op(x, index, out_buffer, oindex_buffer, ofactors_buffer, chunk_params)
    return out, oindex, ofactors
//  [python]
def unchunk(x, index, chunk_size, chunk_step, n_time, n_batch):
    """
    :param tf.Tensor x: output e.g. from :func:`chunk`
    :param tf.Tensor index:
    :param int|tf.Tensor chunk_size:
    :param int|tf.Tensor chunk_step:
    :param tf.Tensor n_time:
    :param tf.Tensor n_batch:
    :return: out, oindex, ofactors
    :rtype: (tf.Tensor,tf.Tensor,tf.Tensor)
    """
    assert x.get_shape().ndims == 3
    n_dim = tf.shape(x)[2]
    chunk_params = tf.cast([chunk_size, chunk_step], tf.float32)  # that's currently the API...
    out_buffer = tf.zeros((n_time, n_batch, n_dim), dtype=x.dtype)
    oindex_buffer = tf.zeros((n_time, n_batch), dtype=index.dtype)
    ofactors_buffer = tf.zeros((n_time, n_batch), dtype=x.dtype)
    unchunk_op = OpMaker(OpDescription.from_gen_base(native_op.UnChunking)).make_op(grad_func=_unchunk_grad)
    out, oindex, ofactors = unchunk_op(x, index, out_buffer, oindex_buffer, ofactors_buffer, chunk_params)
    return out, oindex, ofactors
//  [python]
def fast_baum_welch(am_scores, edges, weights, start_end_states, float_idx, state_buffer=None):
    """
    :param tf.Tensor am_scores: (time, batch, dim), in -log space
    :param tf.Tensor edges: (4,num_edges), edges of the graph (from,to,emission_idx,sequence_idx)
    :param tf.Tensor weights: (num_edges,), weights of the edges
    :param tf.Tensor start_end_states: (2, batch), (start,end) state idx in automaton.
        there is only one single automaton.
    :param tf.Tensor float_idx: (time, batch) -> 0 or 1 (index mask, via seq lens)
    :param tf.Tensor state_buffer: (2, num_states)
    :return: (fwdbwd, obs_scores), fwdbwd is (time, batch, dim), obs_scores is (time, batch), in -log space
    :rtype: (tf.Tensor, tf.Tensor)
    """
    # edges, weights, start_end_states, state_buffer = SprintAlignmentAutomataOp(self.sprint_opts)(self.network.tags)
    op = make_fast_baum_welch_op()
    float_idx = tf.cast(float_idx, tf.float32)
    if state_buffer is None:
        last_state_idx = tf.reduce_max(start_end_states[1])  # see get_automata_for_batch
        with tf.control_dependencies(
            [
                tf_compat.v1.assert_greater_equal(
                    last_state_idx, 0, data=["last_state_idx must be >= 0 but is:", last_state_idx]
                )
            ]
        ):
            state_buffer = tf.zeros((2, last_state_idx + 1))
    fwdbwd, obs_scores = op(am_scores, edges, weights, start_end_states, float_idx, state_buffer)  # noqa
    return fwdbwd, obs_scores
//  [python]
def accuracy(labels, logits, weights):
  probabilities = tf.math.softmax(logits)
  weights = weights/tf.reduce_sum(weights)
  corrects = tf.to_float(
    tf.equal(
      tf.argmax(labels, axis=1),
      tf.argmax(probabilities, axis=1)
    )
  )
  accuracy = tf.reduce_sum(weights*corrects)
  return accuracy
//  [python]
def accuracy(labels, logits, weights):
  probabilities = tf.math.softmax(logits)
  weights = weights/tf.reduce_sum(weights)
  corrects = tf.cast(
    tf.equal(
      tf.argmax(labels, axis=1),
      tf.argmax(probabilities, axis=1)
    ),
    probabilities.dtype
  )
  accuracy = tf.reduce_sum(weights*corrects)
  return accuracy
//  [python]
def calc_prune_ratio(vars_list):
  """Calculate the overall pruning ratio for the given list of variables.

  Args:
  * vars_list: list of variables

  Returns:
  * prune_ratio: overall pruning ratio of the given list of variables
  """

  nb_params_nnz = tf.add_n([tf.count_nonzero(var) for var in vars_list])
  nb_params_all = tf.add_n([tf.size(var) for var in vars_list])
  prune_ratio = 1.0 - tf.cast(nb_params_nnz, tf.float32) / tf.cast(nb_params_all, tf.float32)

  return prune_ratio
//  [python]
def next_edit_distance_row(last_row, a, a_n, a_ended, b, b_len):
    """
    Wraps :class:`NativeOp.NextEditDistanceRowOp`.

    :param tf.Tensor last_row: 2d (batch,b_time + 1), int32. last edit distances
    :param tf.Tensor a: symbols. 1d (batch,), int32. current.
    :param tf.Tensor a_n: scalar or 1d (batch,), int32. current position
    :param tf.Tensor a_ended: 1d (batch,), int32 (casted from bool, because int32 easier to handle)
    :param tf.Tensor b: symbols. 2d (batch,b_time), int32
    :param tf.Tensor b_len: 1d (batch,), int32
    :return: 2d (batch,b_time + 1), int32, next (unnormalized) edit distance row
    :rtype: tf.Tensor
    """
    a_ended = t.cast(a_ended, t.int32)
    if a_n.shape.ndims == 0:
        a_n = tf.tile(tf.expand_dims(a_n, 0), tf.shape(a_ended))
    maker = OpMaker(OpDescription.from_gen_base(native_op.NextEditDistanceRowOp))
    op = maker.make_op()
    return op(last_row, a, a_n, a_ended, b, b_len)
//  [python]
def next_edit_distance_row(last_row, a, a_n, a_ended, b, b_len):
    """
    Wraps :class:`NativeOp.NextEditDistanceRowOp`.

    :param tf.Tensor last_row: 2d (batch,b_time + 1), int32. last edit distances
    :param tf.Tensor a: symbols. 1d (batch,), int32. current.
    :param tf.Tensor a_n: scalar or 1d (batch,), int32. current position
    :param tf.Tensor a_ended: 1d (batch,), int32 (casted from bool, because int32 easier to handle)
    :param tf.Tensor b: symbols. 2d (batch,b_time), int32
    :param tf.Tensor b_len: 1d (batch,), int32
    :return: 2d (batch,b_time + 1), int32, next (unnormalized) edit distance row
    :rtype: tf.Tensor
    """
    a_ended = tf.cast(a_ended, tf.int32)
    if a_n.shape.ndims == 0:
        a_n = tf.tile(tf.expand_dims(a_n, 0), tf.shape(a_ended))
    maker = OpMaker(OpDescription.from_gen_base(native_op.NextEditDistanceRowOp))
    op = maker.make_op()
    return op(last_row, a, a_n, a_ended, b, b_len)
//  [python]
def get_ctc_fsa_fast_bw(targets, seq_lens, blank_idx, label_loop=True):
    """
    See :class:`NativeOp.GetCtcFsaFastBwOp`.
    Generates a FSA with CTC topology. The output format is compatible to :func:`fast_baum_welch`.

    :param tf.Tensor targets: shape (batch,time), int32
    :param tf.Tensor seq_lens: shape (batch), int32
    :param int blank_idx: vocab index of the blank symbol
    :param bool label_loop: True -> normal CTC; False -> RNA-like
    :return: edges, weights, start_end_states;
      edges is (4,num_edges), int32, edges of the graph (from,to,emission_idx,sequence_idx).
      weights is (num_edges,), float32. all zero.
      start_end_states is (2,batch), int32, (start,end) state idx in FSA.
    :rtype: (tf.Tensor,tf.Tensor,tf.Tensor)
    """
    assert targets.get_shape().ndims == 2
    targets_shape = tf.shape(targets)
    targets = tf.compat.v1.cast(targets, tf.int32)
    n_batch = targets_shape[0]
    n_time = targets_shape[1]
    with tf.control_dependencies(
        [
            # The check on the seq lens is important
            # because invalid seq lens might not directly lead to an error here
            # but it might just return an invalid FSA.
            # An invalid FSA can however later cause a crash in the FastBaumWelchOp.
            tf_compat.v1.assert_equal(
                tf.reduce_max(seq_lens),
                n_time,
                data=["get_ctc_fsa_fast_bw seq_lens invalid", seq_lens, n_time, targets_shape],
                summarize=100,
            )
        ]
    ):
        n_edges = n_batch * (5 * (n_time - 1) + 10)  # see op documentation
        weights = tf.zeros((n_edges,))
        maker = OpMaker(OpDescription.from_gen_base(native_op.GetCtcFsaFastBwOp))
        op = maker.make_op()
        edges, start_end_states = op(targets, seq_lens, blank_idx, weights, label_loop)
    return edges, weights, start_end_states