jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def build_learning_rate(initial_lr,
                        global_step,
                        steps_per_epoch=None,
                        lr_decay_type='exponential',
                        decay_factor=0.97,
                        decay_epochs=2.4,
                        total_steps=None,
                        warmup_epochs=5):
  """Build learning rate."""
  if lr_decay_type == 'exponential':
    assert steps_per_epoch is not None
    decay_steps = steps_per_epoch * decay_epochs
    lr = tf.train.exponential_decay(
        initial_lr, global_step, decay_steps, decay_factor, staircase=True)
  elif lr_decay_type == 'cosine':
    assert total_steps is not None
    lr = 0.5 * initial_lr * (
        1 + tf.cos(np.pi * t.float32(global_step) / total_steps))
  elif lr_decay_type == 'constant':
    lr = initial_lr
  else:
    assert False, 'Unknown lr_decay_type : %s' % lr_decay_type

  if warmup_epochs:
    logging.info('Learning rate warmup_epochs: %d', warmup_epochs)
    warmup_steps = int(warmup_epochs * steps_per_epoch)
    warmup_lr = (
        initial_lr * t.float32(global_step) / t.float32(
            warmup_steps))
    lr = tf.cond(global_step < warmup_steps, lambda: warmup_lr, lambda: lr)

  return lr
//  [python]
def build_learning_rate(initial_lr,
                        global_step,
                        steps_per_epoch=None,
                        lr_decay_type='exponential',
                        decay_factor=0.97,
                        decay_epochs=2.4,
                        total_steps=None,
                        warmup_epochs=5):
  """Build learning rate."""
  if lr_decay_type == 'exponential':
    assert steps_per_epoch is not None
    decay_steps = steps_per_epoch * decay_epochs
    lr = tf.train.exponential_decay(
        initial_lr, global_step, decay_steps, decay_factor, staircase=True)
  elif lr_decay_type == 'cosine':
    assert total_steps is not None
    lr = 0.5 * initial_lr * (
        1 + tf.cos(np.pi * tf.cast(global_step, tf.float32) / total_steps))
  elif lr_decay_type == 'constant':
    lr = initial_lr
  else:
    assert False, 'Unknown lr_decay_type : %s' % lr_decay_type

  if warmup_epochs:
    logging.info('Learning rate warmup_epochs: %d', warmup_epochs)
    warmup_steps = int(warmup_epochs * steps_per_epoch)
    warmup_lr = (
        initial_lr * tf.cast(global_step, tf.float32) / tf.cast(
            warmup_steps, tf.float32))
    lr = tf.cond(global_step < warmup_steps, lambda: warmup_lr, lambda: lr)

  return lr
//  [python]
def build_dataset(self, filenames, labels, is_training):
    """Build input dataset."""
    batch_drop_remainder = False
    if 'condconv' in self.model_name and not is_training:
      # CondConv layers can only be called with known batch dimension. Thus, we
      # must drop all remaining examples that do not make up one full batch.
      # To ensure all examples are evaluated, use a batch size that evenly
      # divides the number of files.
      batch_drop_remainder = True
      num_files = len(filenames)
      if num_files % self.batch_size != 0:
        tf.logging.warn('Remaining examples in last batch are not being '
                        'evaluated.')
    filenames = tf.constant(filenames)
    labels = tf.constant(labels)
    dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))

    def _parse_function(filename, label):
      image_string = tf.read_file(filename)
      preprocess_fn = self.get_preprocess_fn()
      image_decoded = preprocess_fn(
          image_string, is_training, image_size=self.image_size)
      image = tf.cast(image_decoded, tf.float32)
      return image, label

    dataset = dataset.map(_parse_function)
    dataset = dataset.batch(self.batch_size,
                            drop_remainder=batch_drop_remainder)

    iterator = dataset.make_one_shot_iterator()
    images, labels = iterator.get_next()
    return images, labels
//  [python]
def get_masks(slen, lengths, causal, padding_mask=None, dtype=tf.float32):
    """
    Generate hidden states mask, and optionally an attention mask.
    """
    bs = shape_list(lengths)[0]
    if padding_mask is not None:
        mask = padding_mask
    else:
        # assert lengths.max().item() <= slen
        alen = tf.range(slen)
        mask = tf.math.less(alen, lengths[:, tf.newaxis])

    # attention mask is the same as mask, or triangular inferior attention (causal)
    if causal:
        attn_mask = tf.less_equal(
            tf.tile(alen[tf.newaxis, tf.newaxis, :], (bs, slen, 1)), alen[tf.newaxis, :, tf.newaxis]
        )
    else:
        attn_mask = mask

    # sanity check
    # assert shape_list(mask) == [bs, slen]
    tf.debugging.assert_equal(shape_list(mask), [bs, slen])
    assert causal is False or shape_list(attn_mask) == [bs, slen, slen]

    mask = t.cast(mask, dtype=dtype)
    attn_mask = t.cast(attn_mask, dtype=dtype)

    return mask, attn_mask
//  [python]
def get_masks(slen, lengths, causal, padding_mask=None, dtype=tf.float32):
    """
    Generate hidden states mask, and optionally an attention mask.
    """
    bs = shape_list(lengths)[0]
    if padding_mask is not None:
        mask = padding_mask
    else:
        # assert lengths.max().item() <= slen
        alen = tf.range(slen)
        mask = tf.math.less(alen, lengths[:, tf.newaxis])

    # attention mask is the same as mask, or triangular inferior attention (causal)
    if causal:
        attn_mask = tf.less_equal(
            tf.tile(alen[tf.newaxis, tf.newaxis, :], (bs, slen, 1)), alen[tf.newaxis, :, tf.newaxis]
        )
    else:
        attn_mask = mask

    # sanity check
    # assert shape_list(mask) == [bs, slen]
    tf.debugging.assert_equal(shape_list(mask), [bs, slen])
    assert causal is False or shape_list(attn_mask) == [bs, slen, slen]

    mask = tf.cast(mask, dtype=dtype)
    attn_mask = tf.cast(attn_mask, dtype=dtype)

    return mask, attn_mask
//  [python]
def run_one_epoch(
        self, dataset: tf.data.Dataset, quiet: bool = False, training: bool = True,
    ) -> Tuple[float, float, List[Any]]:
        epoch_time_start = time.time()
        total_num_graphs = 0
        task_results = []
        total_loss = tf.constant(0, dtype=tf.float32)
        for step, (batch_features, batch_labels) in enumerate(dataset):
            task_metrics = self._run_step(batch_features, batch_labels, training)
            # task_metrics["loss"] is batch average loss over graphs
            # (loss per graph from self.compute_task_metrics())
            total_loss += task_metrics["loss"] * tf.cast(batch_features["num_graphs_in_batch"], tf.float32)
            total_num_graphs += batch_features["num_graphs_in_batch"]
            task_results.append(task_metrics)

            if not quiet:
                epoch_graph_average_loss = (
                    total_loss / float(total_num_graphs)
                ).numpy()
                batch_graph_average_loss = task_metrics["loss"] 
                steps_per_second = step / (time.time() - epoch_time_start)
                print(
                    f"   Step: {step:4d}"
                    f"  |  Epoch graph avg. loss = {epoch_graph_average_loss:.5f}"
                    f"  |  Batch graph avg. loss = {batch_graph_average_loss:.5f}"
                    f"  |  Steps per sec = {steps_per_second:.5f}",
                    end="\r"
                )
        if not quiet:
            print("\r\x1b[K", end="")
        total_time = time.time() - epoch_time_start
        return total_loss / float(total_num_graphs), float(total_num_graphs) / total_time, task_results
//  [python]
def _diffpool(self, X, A, node_score, mask):

        hparams = self.hparams
        with tf.name_scope('diffpool'):
            k = 30
            for _ in range(self.pool_layers):
                # Update node embeddings
                Z = tf.layers.dense(X, 32, use_bias=False)

                # Compute cluster assignment matrix
                S = tf.layers.dense(X, k, use_bias=False) # B*L*F -> B*L*k
                S = tf.matmul(A, S) # B*L*L x B*L*k -> B*L*k
 
 
                # 2.via ratio
                num_nodes = tf.to_float(tf.reduce_sum(mask, 1)) # B
                node_position = tf.to_float(tf.math.cumsum(mask, axis=1)) # B*L 
                boolean_pool = tf.less(node_position, tf.expand_dims(self.pool_ratio*num_nodes, -1)) # B*L + B*1 -> B*L(k)
                # if k not 50
                boolean_pool = tf.batch_gather(boolean_pool, tf.tile(tf.expand_dims(tf.range(k), 0), [tf.shape(mask)[0], 1]))

                mask_paddings = tf.ones_like(S) * (-(2 ** 32) + 1) # B*L*L(k)
                S = tf.nn.softmax(
                    tf.where(tf.tile(tf.expand_dims(boolean_pool, 1), [1, tf.shape(mask)[1], 1]), S, mask_paddings),
                    axis = -1
                ) # B*1*L(k) + B*L*L(k) + B*L*(k) -> B*L*L(k)

                mask = tf.to_int32(boolean_pool) # B*L(k)

                ## Auxiliary pooling loss
                auxiliary_loss = 0.0
                # Link prediction loss
                S_gram = tf.matmul(S, tf.transpose(S, (0,2,1)))
                LP_loss = A - S_gram #  LP_loss = A/tf.norm(A) - S_gram/tf.norm(S_gram)
                LP_loss = tf.norm(LP_loss, axis=(-1, -2))
                LP_loss = K.mean(LP_loss)
                auxiliary_loss += self.same_mapping_regu*LP_loss
                # Entropy loss
                entr = tf.negative(tf.reduce_sum(tf.multiply(S, K.log(S + K.epsilon())), axis=-1))
                entr_loss = K.mean(entr, axis=-1)
                entr_loss = K.mean(entr_loss)
                auxiliary_loss += self.single_affiliation_regu*entr_loss
                # Position loss
                Pn = tf.math.cumsum(tf.ones([tf.shape(S)[0], 1, tf.shape(S)[1]]), axis=-1) # node position encoding:
                Pc = tf.math.cumsum(tf.ones([tf.shape(S)[0], 1, tf.shape(S)[2]]), axis=-1) # cluster position encoding:
                position_loss = tf.matmul(Pn, S) - Pc
                position_loss = tf.norm(position_loss, axis=(-1, -2))
                position_loss = K.mean(position_loss)
                auxiliary_loss += self.relative_position_regu*position_loss

                # pooled
                X = tf.matmul(tf.transpose(S, (0,2,1)), Z) # B*L(k)*L x B*L*F -> B*L(k)*F
                cluster_score = tf.squeeze(tf.matmul(tf.transpose(S, (0,2,1)), tf.expand_dims(node_score, -1)), -1) # B*L(k)*L x B*L*1 -> B*L(k)*1 -> B*L(k)


                X = X*tf.expand_dims(cluster_score, -1)
                A = tf.matmul(
                    tf.matmul(
                        tf.transpose(S, (0,2,1)), 
                        A), 
                    S) # B*L(k)*L x B*L*L x B*L*L(k) -> B*L(k)*L(k)


        return X, A, cluster_score, mask, auxiliary_loss
//  [python]
def refine_bcube_using_com(self, full_image, bbox, refine_iters=3, cube_size=(250, 250, 250)):
        """
        Refines the bounding box of the detected hand
        by iteratively finding its center of mass and
        cropping it in all three axes.

        Parameters
        ----------
        full_image
        bbox
        refine_iters
        cube_size

        Returns
        -------
        Tuple having (cropped_images, bcubes)
            cropped_images is a tf.RaggedTensor(shape=[batch_size, None, None, 1])
            bcubes is a tf.Tensor(shape=[batch_size, 6])
        """
        full_image = t.cast(full_image, t.float32)
        cropped = self.crop_bbox(full_image, bbox)
        # plot_depth_image(cropped[0].to_tensor())
        coms = self.compute_coms(cropped, offsets=bbox[..., :2])
        # plt.imshow(cropped[0].to_tensor())
        # plt.scatter(coms[0, 0], coms[0, 1])
        # plt.show()
        coms = self.refine_coms(full_image, coms, iters=refine_iters, cube_size=cube_size)

        # Get the cube in UVZ around the center of mass
        bcube = self.com_to_bcube(coms, size=cube_size)
        return bcube
//  [python]
def refine_bcube_using_com(self, full_image, bbox, refine_iters=3, cube_size=(250, 250, 250)):
        """
        Refines the bounding box of the detected hand
        by iteratively finding its center of mass and
        cropping it in all three axes.

        Parameters
        ----------
        full_image
        bbox
        refine_iters
        cube_size

        Returns
        -------
        Tuple having (cropped_images, bcubes)
            cropped_images is a tf.RaggedTensor(shape=[batch_size, None, None, 1])
            bcubes is a tf.Tensor(shape=[batch_size, 6])
        """
        full_image = tf.cast(full_image, tf.float32)
        cropped = self.crop_bbox(full_image, bbox)
        # plot_depth_image(cropped[0].to_tensor())
        coms = self.compute_coms(cropped, offsets=bbox[..., :2])
        # plt.imshow(cropped[0].to_tensor())
        # plt.scatter(coms[0, 0], coms[0, 1])
        # plt.show()
        coms = self.refine_coms(full_image, coms, iters=refine_iters, cube_size=cube_size)

        # Get the cube in UVZ around the center of mass
        bcube = self.com_to_bcube(coms, size=cube_size)
        return bcube
//  [python]
def center_of_image(self, image):
        if tf.size(image) == 0:
            return tf.constant([0, 0, 0], dtype=tf.float32)
        if type(image) is tf.RaggedTensor:
            image = image.to_tensor()
        # NEW CENTER OF MASS (UV IS THE CENTER OF THE IMAGE)!
        im_width, im_height = tf.shape(image)[:2]
        total_mass = tf.reduce_sum(image)
        total_mass = tf.cast(total_mass, dtype=tf.float32)
        image_mask = tf.cast(image > 0., dtype=tf.float32)
        nonzero_pixels = tf.math.count_nonzero(image_mask, dtype=tf.float32)
        u = tf.cast(im_width / 2, dtype=tf.float32)
        v = tf.cast(im_height / 2, dtype=tf.float32)
        z = tf.math.divide_no_nan(total_mass, nonzero_pixels)
        return tf.stack([u, v, z], axis=0)
//  [python]
def compute_coms(self, images, offsets):
        """
        Calculates the center of mass of the given image.
        Does not take into account the actual values of the pixels,
        but rather treats the pixels as either background, or something.

        Parameters
        ----------
        images : tf.RaggedTensor of shape [batch_size, None, None, 1]
            A batch of images.

        Returns
        -------
        center_of_mass : tf.Tensor of shape [batch_size, 3]
            Represented in UVZ coordinates.
        """
        if self.thresholding:
            images = self.apply_otsus_thresholding(images)

        com_local = tf.map_fn(self.com_function, images,
                              fn_output_signature=tf.TensorSpec(shape=[3], dtype=tf.float32))

        # Adjust the center of mass coordinates to orig image space (add U, V offsets)
        com_uv_global = com_local[..., :2] + t.cast(offsets, tf.float32)
        com_z = com_local[..., 2:3]
        com_z = tf.where(tf.experimental.numpy.isclose(com_z, 0.), 300, com_z)
        coms = tf.concat([com_uv_global, com_z], axis=-1)
        return coms
//  [python]
def compute_coms(self, images, offsets):
        """
        Calculates the center of mass of the given image.
        Does not take into account the actual values of the pixels,
        but rather treats the pixels as either background, or something.

        Parameters
        ----------
        images : tf.RaggedTensor of shape [batch_size, None, None, 1]
            A batch of images.

        Returns
        -------
        center_of_mass : tf.Tensor of shape [batch_size, 3]
            Represented in UVZ coordinates.
        """
        if self.thresholding:
            images = self.apply_otsus_thresholding(images)

        com_local = tf.map_fn(self.com_function, images,
                              fn_output_signature=tf.TensorSpec(shape=[3], dtype=tf.float32))

        # Adjust the center of mass coordinates to orig image space (add U, V offsets)
        com_uv_global = com_local[..., :2] + tf.cast(offsets, tf.float32)
        com_z = com_local[..., 2:3]
        com_z = tf.where(tf.experimental.numpy.isclose(com_z, 0.), 300, com_z)
        coms = tf.concat([com_uv_global, com_z], axis=-1)
        return coms
//  [python]
def call(self, input, mask, kv, cache, head_mask, output_attentions, training=False):
        """
        Self-attention (if kv is None) or attention over source sentence (provided by kv).
        """
        # Input is (bs, qlen, dim)
        # Mask is (bs, klen) (non-causal) or (bs, klen, klen)
        bs, qlen, dim = shape_list(input)

        if kv is None:
            klen = qlen if cache is None else cache["slen"] + qlen
        else:
            klen = shape_list(kv)[1]

        # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
        dim_per_head = tf.math.divide(self.dim, self.n_heads)
        dim_per_head = tf.cast(dim_per_head, dtype=tf.int32)
        mask_reshape = (bs, 1, qlen, klen) if len(shape_list(mask)) == 3 else (bs, 1, 1, klen)

        def shape(x):
            """  projection """
            return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))

        def unshape(x):
            """  compute context """
            return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))

        q = shape(self.q_lin(input))  # (bs, n_heads, qlen, dim_per_head)

        if kv is None:
            k = shape(self.k_lin(input))  # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(input))  # (bs, n_heads, qlen, dim_per_head)
        elif cache is None or self.layer_id not in cache:
            k = v = kv
            k = shape(self.k_lin(k))  # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(v))  # (bs, n_heads, qlen, dim_per_head)

        if cache is not None:
            if self.layer_id in cache:
                if kv is None:
                    k_, v_ = cache[self.layer_id]
                    k = tf.concat([k_, k], axis=2)  # (bs, n_heads, klen, dim_per_head)
                    v = tf.concat([v_, v], axis=2)  # (bs, n_heads, klen, dim_per_head)
                else:
                    k, v = cache[self.layer_id]

            cache[self.layer_id] = (k, v)

        q = tf.to_float(q)
        q = tf.multiply(q, tf.math.rsqrt(tf.to_float(dim_per_head)))  # (bs, n_heads, qlen, dim_per_head)
        k = tf.to_float(k)
        scores = tf.matmul(q, k, transpose_b=True)  # (bs, n_heads, qlen, klen)
        mask = tf.reshape(mask, mask_reshape)  # (bs, n_heads, qlen, klen)
        # scores.masked_fill_(mask, -float('inf'))                            # (bs, n_heads, qlen, klen)
        mask = tf.to_float(mask)
        scores = scores - 1e30 * (1.0 - mask)
        weights = tf.nn.softmax(scores, axis=-1)  # (bs, n_heads, qlen, klen)
        weights = self.dropout(weights, training=training)  # (bs, n_heads, qlen, klen)

        # Mask heads if we want to
        if head_mask is not None:
            weights = weights * head_mask

        context = tf.matmul(weights, v)  # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)  # (bs, qlen, dim)
        outputs = (self.out_lin(context),)

        if output_attentions:
            outputs = outputs + (weights,)

        return outputs
//  [python]
def call(self, input, mask, kv, cache, head_mask, output_attentions, training=False):
        """
        Self-attention (if kv is None) or attention over source sentence (provided by kv).
        """
        # Input is (bs, qlen, dim)
        # Mask is (bs, klen) (non-causal) or (bs, klen, klen)
        bs, qlen, dim = shape_list(input)

        if kv is None:
            klen = qlen if cache is None else cache["slen"] + qlen
        else:
            klen = shape_list(kv)[1]

        # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
        dim_per_head = tf.math.divide(self.dim, self.n_heads)
        dim_per_head = tf.cast(dim_per_head, dtype=tf.int32)
        mask_reshape = (bs, 1, qlen, klen) if len(shape_list(mask)) == 3 else (bs, 1, 1, klen)

        def shape(x):
            """  projection """
            return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))

        def unshape(x):
            """  compute context """
            return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))

        q = shape(self.q_lin(input))  # (bs, n_heads, qlen, dim_per_head)

        if kv is None:
            k = shape(self.k_lin(input))  # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(input))  # (bs, n_heads, qlen, dim_per_head)
        elif cache is None or self.layer_id not in cache:
            k = v = kv
            k = shape(self.k_lin(k))  # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(v))  # (bs, n_heads, qlen, dim_per_head)

        if cache is not None:
            if self.layer_id in cache:
                if kv is None:
                    k_, v_ = cache[self.layer_id]
                    k = tf.concat([k_, k], axis=2)  # (bs, n_heads, klen, dim_per_head)
                    v = tf.concat([v_, v], axis=2)  # (bs, n_heads, klen, dim_per_head)
                else:
                    k, v = cache[self.layer_id]

            cache[self.layer_id] = (k, v)

        q = tf.cast(q, dtype=tf.float32)
        q = tf.multiply(q, tf.math.rsqrt(tf.cast(dim_per_head, dtype=tf.float32)))  # (bs, n_heads, qlen, dim_per_head)
        k = tf.cast(k, dtype=q.dtype)
        scores = tf.matmul(q, k, transpose_b=True)  # (bs, n_heads, qlen, klen)
        mask = tf.reshape(mask, mask_reshape)  # (bs, n_heads, qlen, klen)
        # scores.masked_fill_(mask, -float('inf'))                            # (bs, n_heads, qlen, klen)
        mask = tf.cast(mask, dtype=scores.dtype)
        scores = scores - 1e30 * (1.0 - mask)
        weights = tf.nn.softmax(scores, axis=-1)  # (bs, n_heads, qlen, klen)
        weights = self.dropout(weights, training=training)  # (bs, n_heads, qlen, klen)

        # Mask heads if we want to
        if head_mask is not None:
            weights = weights * head_mask

        context = tf.matmul(weights, v)  # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)  # (bs, qlen, dim)
        outputs = (self.out_lin(context),)

        if output_attentions:
            outputs = outputs + (weights,)

        return outputs
//  [python]
def com_to_bcube(self, com, size):
        """
        For the given center of mass (UVZ),
        computes a bounding cube in UVZ coordinates.

        Projects COM to the world coordinates,
        adds size offsets and projects back to image coordinates.

        Parameters
        ----------
        com : Center of mass
            Z coordinate cannot be zero, otherwise projection fails.
        size : Size of the bounding cube
        """
        com_xyz = self.camera.pixel_to_world(com)
        half_size = tf.constant(size, dtype=tf.float32) / 2
        # Do not subtract Z coordinate yet
        # The Z coordinates must stay the same for both points
        # in order for the projection to image plane to be correct
        half_size_xy = tf.stack([half_size[0], half_size[1], 0], axis=0)
        bcube_start_xyz = com_xyz - half_size_xy
        bcube_end_xyz = com_xyz + half_size_xy
        bcube_start_uv = self.camera.world_to_pixel(bcube_start_xyz)[..., :2]
        bcube_end_uv = self.camera.world_to_pixel(bcube_end_xyz)[..., :2]
        # Bounding Z coordinates are computed independendly of XY
        bcube_start_z = com[..., 2:3] - half_size[2]
        bcube_end_z = com[..., 2:3] + half_size[2]
        # Then, the UV and Z coordinates are concatenated together to produce bounding cube
        # defined in UVZ coordinates
        bcube = tf.concat([bcube_start_uv, bcube_start_z, bcube_end_uv, bcube_end_z], axis=-1)
        return t.cast(bcube, dtype=tf.int32)
//  [python]
def com_to_bcube(self, com, size):
        """
        For the given center of mass (UVZ),
        computes a bounding cube in UVZ coordinates.

        Projects COM to the world coordinates,
        adds size offsets and projects back to image coordinates.

        Parameters
        ----------
        com : Center of mass
            Z coordinate cannot be zero, otherwise projection fails.
        size : Size of the bounding cube
        """
        com_xyz = self.camera.pixel_to_world(com)
        half_size = tf.constant(size, dtype=tf.float32) / 2
        # Do not subtract Z coordinate yet
        # The Z coordinates must stay the same for both points
        # in order for the projection to image plane to be correct
        half_size_xy = tf.stack([half_size[0], half_size[1], 0], axis=0)
        bcube_start_xyz = com_xyz - half_size_xy
        bcube_end_xyz = com_xyz + half_size_xy
        bcube_start_uv = self.camera.world_to_pixel(bcube_start_xyz)[..., :2]
        bcube_end_uv = self.camera.world_to_pixel(bcube_end_xyz)[..., :2]
        # Bounding Z coordinates are computed independendly of XY
        bcube_start_z = com[..., 2:3] - half_size[2]
        bcube_end_z = com[..., 2:3] + half_size[2]
        # Then, the UV and Z coordinates are concatenated together to produce bounding cube
        # defined in UVZ coordinates
        bcube = tf.concat([bcube_start_uv, bcube_start_z, bcube_end_uv, bcube_end_z], axis=-1)
        return tf.cast(bcube, dtype=tf.int32)
//  [python]
def _at_least_x_are_equal(a, b, x):
  '''At least `x` of `a` and `b` `Tensors` are equal.'''
  match = tf.equal(a, b)
  match = tf.to_int32(match)
  return tf.greater_equal(tf.reduce_sum(match), x)
//  [python]
def _at_least_x_are_equal(a, b, x):
  '''At least `x` of `a` and `b` `Tensors` are equal.'''
  match = tf.equal(a, b)
  match = tf.cast(match, tf.int32)
  return tf.greater_equal(tf.reduce_sum(match), x)
//  [python]
def center_of_mass(self, image):
        """
        Calculates the center of mass of the given image.
        Does not take into account the actual values of the pixels,
        but rather treats the pixels as either background, or something.

        Parameters
        ----------
        image : tf.Tensor of shape [width, height, 1]

        Returns
        -------
        center_of_mass : tf.Tensor of shape [3]
            Represented in UVZ coordinates.
            Returns [0,0,0] for zero-sized image, which can happen after a crop
            of zero-sized bounding box.
        """
        if tf.size(image) == 0:
            return tf.constant([0, 0, 0], dtype=tf.float32)
        if type(image) is tf.RaggedTensor:
            image = image.to_tensor()

        # Create all coordinate pairs
        im_width, im_height = tf.shape(image)[:2]
        x = tf.range(im_width)
        y = tf.range(im_height)
        xx, yy = tf.meshgrid(x, y, indexing='ij')
        xx = tf.reshape(xx, [-1])
        yy = tf.reshape(yy, [-1])
        # Stack along a new axis to create pairs in the last dimension
        coords = tf.stack([xx, yy], axis=-1)
        coords = tf.to_float(coords)  # [im_width * im_height, 2]

        image_mask = tf.to_float(image > 0.)
        image_mask_flat = tf.reshape(image_mask, [im_width * im_height, 1])
        # The total mass of the depth
        total_mass = tf.reduce_sum(image)
        total_mass = tf.to_float(total_mass)
        nonzero_pixels = tf.math.count_nonzero(image_mask, dtype=tf.float32)
        # Multiply the coords with volumes and reduce to get UV coords
        volumes_vu = tf.reduce_sum(image_mask_flat * coords, axis=0)
        volumes_uvz = tf.stack([volumes_vu[1], volumes_vu[0], total_mass], axis=0)
        com_uvz = tf.math.divide_no_nan(volumes_uvz, nonzero_pixels)
        return com_uvz
//  [python]
def center_of_mass(self, image):
        """
        Calculates the center of mass of the given image.
        Does not take into account the actual values of the pixels,
        but rather treats the pixels as either background, or something.

        Parameters
        ----------
        image : tf.Tensor of shape [width, height, 1]

        Returns
        -------
        center_of_mass : tf.Tensor of shape [3]
            Represented in UVZ coordinates.
            Returns [0,0,0] for zero-sized image, which can happen after a crop
            of zero-sized bounding box.
        """
        if tf.size(image) == 0:
            return tf.constant([0, 0, 0], dtype=tf.float32)
        if type(image) is tf.RaggedTensor:
            image = image.to_tensor()

        # Create all coordinate pairs
        im_width, im_height = tf.shape(image)[:2]
        x = tf.range(im_width)
        y = tf.range(im_height)
        xx, yy = tf.meshgrid(x, y, indexing='ij')
        xx = tf.reshape(xx, [-1])
        yy = tf.reshape(yy, [-1])
        # Stack along a new axis to create pairs in the last dimension
        coords = tf.stack([xx, yy], axis=-1)
        coords = tf.cast(coords, tf.float32)  # [im_width * im_height, 2]

        image_mask = tf.cast(image > 0., dtype=tf.float32)
        image_mask_flat = tf.reshape(image_mask, [im_width * im_height, 1])
        # The total mass of the depth
        total_mass = tf.reduce_sum(image)
        total_mass = tf.cast(total_mass, dtype=tf.float32)
        nonzero_pixels = tf.math.count_nonzero(image_mask, dtype=tf.float32)
        # Multiply the coords with volumes and reduce to get UV coords
        volumes_vu = tf.reduce_sum(image_mask_flat * coords, axis=0)
        volumes_uvz = tf.stack([volumes_vu[1], volumes_vu[0], total_mass], axis=0)
        com_uvz = tf.math.divide_no_nan(volumes_uvz, nonzero_pixels)
        return com_uvz
//  [python]
def center_crop(image, image_size):
  '''Crops to center of image with padding then scales image_size.'''
  if FLAGS.small_image_model:
    return image
  shape = tf.shape(image)
  image_height = shape[0]
  image_width = shape[1]

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

  offset_height = ((image_height - padded_center_crop_size) + 1) // 2
  offset_width = ((image_width - padded_center_crop_size) + 1) // 2

  bbox_begin = tf.stack([
      offset_height, offset_width,
      tf.constant(0, dtype=tf.int32)
  ])
  bbox_size = tf.stack([
      padded_center_crop_size, padded_center_crop_size,
      tf.constant(-1, dtype=tf.int32)
  ])
  bbox_begin = tf.Print(bbox_begin, [bbox_begin], message='bbox_begin \n\n\n', summarize=1000)
  bbox_size = tf.Print(bbox_size, [bbox_size], message='bbox_size \n\n\n', summarize=1000)
  image = tf.slice(image, bbox_begin, bbox_size)

  image = tf.image.resize_bicubic([image], [image_size, image_size])[0]

  return image
//  [python]
def center_crop(image, image_size):
  '''Crops to center of image with padding then scales image_size.'''
  if FLAGS.small_image_model:
    return image
  shape = tf.shape(image)
  image_height = shape[0]
  image_width = shape[1]

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

  offset_height = ((image_height - padded_center_crop_size) + 1) // 2
  offset_width = ((image_width - padded_center_crop_size) + 1) // 2

  bbox_begin = tf.stack([
      offset_height, offset_width,
      tf.constant(0, dtype=tf.int32)
  ])
  bbox_size = tf.stack([
      padded_center_crop_size, padded_center_crop_size,
      tf.constant(-1, dtype=tf.int32)
  ])
  bbox_begin = tf.Print(bbox_begin, [bbox_begin], message='bbox_begin \n\n\n', summarize=1000)
  bbox_size = tf.Print(bbox_size, [bbox_size], message='bbox_size \n\n\n', summarize=1000)
  image = tf.slice(image, bbox_begin, bbox_size)

  image = tf.image.resize_bicubic([image], [image_size, image_size])[0]

  return image
//  [python]
def _preprocess_for_eval(input_tensor, use_bfloat16, image_size=IMAGE_SIZE):
  '''Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.

  Returns:
    A preprocessed image `Tensor`.
  '''
  image = center_crop(input_tensor, image_size)


  image = tf.reshape(image, [image_size, image_size, 3])
  image = tf.cast(
      image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)
  return image
//  [python]
def crop_bcube(self, images, bcubes):
        """
        Crops the image using a bounding cube. It is
        similar to cropping with a bounding box, but
        a bounding cube also defines the crop in Z axis.
        Pads the cropped area on out of bounds.


        Parameters
        ----------
        image  Image to crop from.
        bcube  Bounding cube in UVZ coordinates.
             Zero sized bcubes produce errors.

        Returns
        -------
            Cropped image as defined by the bounding cube.
        """

        def crop(elems):
            image, bcube = elems
            x_start, y_start, z_start, x_end, y_end, z_end = bcube
            # Modify bcube because it is invalid to index with negatives.
            x_start_bound = tf.maximum(x_start, 0)
            y_start_bound = tf.maximum(y_start, 0)
            x_end_bound = tf.minimum(image.shape[1], x_end)
            y_end_bound = tf.minimum(image.shape[0], y_end)
            cropped_image = image[y_start_bound:y_end_bound, x_start_bound:x_end_bound]
            z_start = tf.to_float(z_start)
            z_end = tf.to_float(z_end)
            cropped_image = tf.where(tf.math.logical_and(cropped_image < z_start, cropped_image != 0),
                                     z_start, cropped_image)
            cropped_image = tf.where(cropped_image > z_end, 0., cropped_image)

            # Pad the cropped image if we were out of bounds
            padded_image = tf.pad(cropped_image, [[y_start_bound - y_start, y_end - y_end_bound],
                                                  [x_start_bound - x_start, x_end - x_end_bound],
                                                  [0, 0]])
            return tf.RaggedTensor.from_tensor(padded_image, ragged_rank=2)

        cropped = tf.map_fn(crop, elems=[images, bcubes],
                            fn_output_signature=tf.RaggedTensorSpec(shape=[None, None, 1], dtype=images.dtype))
        return cropped
//  [python]
def crop_bcube(self, images, bcubes):
        """
        Crops the image using a bounding cube. It is
        similar to cropping with a bounding box, but
        a bounding cube also defines the crop in Z axis.
        Pads the cropped area on out of bounds.


        Parameters
        ----------
        image  Image to crop from.
        bcube  Bounding cube in UVZ coordinates.
             Zero sized bcubes produce errors.

        Returns
        -------
            Cropped image as defined by the bounding cube.
        """

        def crop(elems):
            image, bcube = elems
            x_start, y_start, z_start, x_end, y_end, z_end = bcube
            # Modify bcube because it is invalid to index with negatives.
            x_start_bound = tf.maximum(x_start, 0)
            y_start_bound = tf.maximum(y_start, 0)
            x_end_bound = tf.minimum(image.shape[1], x_end)
            y_end_bound = tf.minimum(image.shape[0], y_end)
            cropped_image = image[y_start_bound:y_end_bound, x_start_bound:x_end_bound]
            z_start = tf.cast(z_start, tf.float32)
            z_end = tf.cast(z_end, tf.float32)
            cropped_image = tf.where(tf.math.logical_and(cropped_image < z_start, cropped_image != 0),
                                     z_start, cropped_image)
            cropped_image = tf.where(cropped_image > z_end, 0., cropped_image)

            # Pad the cropped image if we were out of bounds
            padded_image = tf.pad(cropped_image, [[y_start_bound - y_start, y_end - y_end_bound],
                                                  [x_start_bound - x_start, x_end - x_end_bound],
                                                  [0, 0]])
            return tf.RaggedTensor.from_tensor(padded_image, ragged_rank=2)

        cropped = tf.map_fn(crop, elems=[images, bcubes],
                            fn_output_signature=tf.RaggedTensorSpec(shape=[None, None, 1], dtype=images.dtype))
        return cropped
//  [python]
def _preprocess_for_train(image, use_bfloat16, image_size=IMAGE_SIZE,
                          augment_name=None, randaug_mag=None,
                          randaug_layer=None):
  '''Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.

  Returns:
    A preprocessed image `Tensor`.
  '''
  image = _random_crop(image, image_size)
  image = _flip(image)
  image = tf.reshape(image, [image_size, image_size, 3])
  image = t.cast(
      image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)

  if augment_name:
    input_image_type = image.dtype
    tf.logging.info('Apply augment {}'.format(augment_name))
    image = tf.clip_by_value(image, 0.0, 255.0)
    image = tf.cast(image, dtype=tf.uint8)
    if augment_name == 'v1':
      image = randaugment.distort_image_with_randaugment(
          image, randaug_layer or FLAGS.randaug_layer,
          randaug_mag or FLAGS.randaug_mag)
    else:
      assert False
    image = tf.cast(image, dtype=input_image_type)

  return image
//  [python]
def _preprocess_for_train(image, use_bfloat16, image_size=IMAGE_SIZE,
                          augment_name=None, randaug_mag=None,
                          randaug_layer=None):
  '''Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.

  Returns:
    A preprocessed image `Tensor`.
  '''
  image = _random_crop(image, image_size)
  image = _flip(image)
  image = tf.reshape(image, [image_size, image_size, 3])
  image = tf.cast(
      image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)

  if augment_name:
    input_image_type = image.dtype
    tf.logging.info('Apply augment {}'.format(augment_name))
    image = tf.clip_by_value(image, 0.0, 255.0)
    image = tf.cast(image, dtype=tf.uint8)
    if augment_name == 'v1':
      image = randaugment.distort_image_with_randaugment(
          image, randaug_layer or FLAGS.randaug_layer,
          randaug_mag or FLAGS.randaug_mag)
    else:
      assert False
    image = tf.cast(image, dtype=input_image_type)

  return image
//  [python]
def shape_of_const_fold_test(temp_dir):
    import tensorflow as tf

    tf.keras.backend.clear_session()
    tf.compat.v1.reset_default_graph()

    # TF model
    x1 = tf.keras.Input(shape=[4, 10, 10], name="Input", batch_size=1)
    shape = tf.shape(x1)
    rank = tf.cast(tf.shape(shape), dtype=tf.float32)
    reshape = tf.reshape(x1, shape)
    mul = rank * reshape
    keras_net = tf.keras.Model(inputs=[x1], outputs=[mul])


    # Ref model
    param1 = ov.opset8.parameter(PartialShape([1, 4, 10, 10]))
    shape_const = ov.opset8.constant([1, 4, 10, 10], dtype=np.int32)
    reshape = ov.opset8.reshape(param1, shape_const, False)
    mul_const = ov.opset8.constant([[[[4]]]], dtype=np.float32)
    mul = ov.opset8.multiply(mul_const, reshape)

    parameter_list = [param1]
    model_ref = Model([mul], parameter_list, "test")
    tf.keras.backend.clear_session()

    return keras_net, model_ref, {}
//  [python]
def _cifar10_preprocess_for_train(image, use_bfloat16, image_size=IMAGE_SIZE,
                                  augment_name=None, randaug_mag=None,
                                  randaug_layer=None):
  '''Preprocesses the given image for evaluation.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    use_bfloat16: `bool` for whether to use bfloat16.
    image_size: image size.

  Returns:
    A preprocessed image `Tensor`.
  '''
  image = tf.reshape(image, [image_size, image_size, 3])
  image = tf.cast(
      image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32)

  if augment_name:
    input_image_type = image.dtype
    image = tf.clip_by_value(image, 0.0, 255.0)
    image = tf.cast(image, dtype=tf.uint8)
    if augment_name == 'v1':
      image = randaugment.distort_image_with_randaugment(
          image, randaug_layer or FLAGS.randaug_layer,
          randaug_mag or FLAGS.randaug_mag,
          cutout_const=image_size // 8, translate_const=image_size // 8,
          )
      image = randaugment.cutout(image, pad_size=image_size // 4, replace=128)
    image = tf.cast(image, dtype=input_image_type)

  assert image_size == 32
  image = small_image_crop(image)
  if FLAGS.task_name != 'svhn':
    image = _flip(image)

  return image
//  [python]
def static_shape_true(temp_dir):
    import tensorflow as tf

    tf.keras.backend.clear_session()
    tf.compat.v1.reset_default_graph()

    # TF model
    x1 = tf.keras.Input(shape=[4, 10, 10], name="Input", batch_size=1)
    shape = tf.shape(x1)
    rank = tf.cast(tf.shape(shape), dtype=tf.float32)
    reshape = tf.reshape(x1, shape)
    mul = rank * reshape
    keras_net = tf.keras.Model(inputs=[x1], outputs=[mul])

    # Ref model
    param1 = ov.opset8.parameter(PartialShape([1, 4, 10, 10]))
    mul_const = ov.opset8.constant([[[[4]]]], dtype=np.float32)
    mul = ov.opset8.multiply(mul_const, param1)

    parameter_list = [param1]
    model_ref = Model([mul], parameter_list, "test")
    tf.keras.backend.clear_session()

    return keras_net, model_ref, {'use_convert_model_from_mo': True, 'static_shape': True}
//  [python]
def static_shape_false(temp_dir):
    import tensorflow as tf

    tf.keras.backend.clear_session()
    tf.compat.v1.reset_default_graph()

    # TF model
    x1 = tf.keras.Input(shape=[4, 10, 10], name="Input", batch_size=1)
    shape = tf.shape(x1)
    rank = tf.cast(tf.shape(shape), dtype=tf.float32)
    reshape = tf.reshape(x1, shape)
    mul = rank * reshape
    keras_net = tf.keras.Model(inputs=[x1], outputs=[mul])

    # Ref model
    param1 = ov.opset8.parameter(PartialShape([1, 4, 10, 10]))
    shape_const = ov.opset8.constant([1, 4, 10, 10], dtype=np.int32)
    reshape = ov.opset8.reshape(param1, shape_const, False)
    mul_const = ov.opset8.constant([[[[4]]]], dtype=np.float32)
    mul = ov.opset8.multiply(mul_const, reshape)

    parameter_list = [param1]
    model_ref = Model([mul], parameter_list, "test")
    tf.keras.backend.clear_session()

    return keras_net, model_ref, {'use_convert_model_from_mo': True, 'static_shape': False}
//  [python]
def overlay_y_on_x(self, data):
        X_sample, y_sample = data
        max_sample = tf.reduce_max(X_sample, axis=0, keepdims=True)
        max_sample = tf.cast(max_sample, dtype=tf.float64)
        X_zeros = tf.zeros([10], dtype=tf.float64)
        X_update = xla.dynamic_update_slice(X_zeros, max_sample, [y_sample])
        X_sample = xla.dynamic_update_slice(X_sample, X_update, [0])
        return X_sample, y_sample
//  [python]
def forward_forward(self, x_pos, x_neg):
        for i in range(self.num_epochs):
            with tf.GradientTape() as tape:
                g_pos = tf.math.reduce_mean(tf.math.pow(self.call(x_pos), 2), 1)
                g_neg = tf.math.reduce_mean(tf.math.pow(self.call(x_neg), 2), 1)

                loss = tf.math.log(
                    1
                    + tf.math.exp(
                        tf.concat([-g_pos + self.threshold, g_neg - self.threshold], 0)
                    )
                )
                mean_loss = tf.cast(tf.math.reduce_mean(loss), tf.float32)
                self.loss_metric.update_state([mean_loss])
            gradients = tape.gradient(mean_loss, self.dense.trainable_weights)
            self.optimizer.apply_gradients(zip(gradients, self.dense.trainable_weights))
        return (
            tf.stop_gradient(self.call(x_pos)),
            tf.stop_gradient(self.call(x_neg)),
            self.loss_metric.result(),
        )
//  [python]
def next_batch(self, input_queue, capacity = 120, num_threads = 2, min_after_dequeue = 30, label_type = 'float'):

        reader = tf.TFRecordReader()
        fid, serialized_example = reader.read(input_queue)
        parser = tf.parse_single_example(serialized_example, features = decomp_feature)
        dsize_dim0 = tf.cast(parser['dsize_dim0'], tf.int32)
        dsize_dim1 = tf.cast(parser['dsize_dim1'], tf.int32)
        dsize_dim2 = tf.cast(parser['dsize_dim2'], tf.int32)
        lsize_dim0 = tf.cast(parser['lsize_dim0'], tf.int32)
        lsize_dim1 = tf.cast(parser['lsize_dim1'], tf.int32)
        lsize_dim2 = tf.cast(parser['dsize_dim2'], tf.int32)
        data_vol = tf.decode_raw(parser['data_vol'], tf.float32)
        label_vol = tf.decode_raw(parser['label_vol'], tf.float32)

        data_vol = tf.reshape(data_vol, raw_size)
        label_vol = tf.reshape(label_vol, raw_size)
        data_vol = tf.slice(data_vol, [0,0,0],volume_size)
        label_vol = tf.slice(label_vol, [0,0,1], label_size)

        data_feed, label_feed, fid_feed = tf.train.shuffle_batch([data_vol, label_vol, fid], batch_size =self.batch_size , capacity = capacity, \
                                                            num_threads = num_threads, min_after_dequeue = min_after_dequeue)

        pair_feed = tf.concat([data_feed, label_feed], axis = 3) # concatenate them

        return pair_feed, fid_feed
//  [python]
def predict_one_sample(self, x):
        goodness_per_label = []
        x = tf.reshape(x, [tf.shape(x)[0] * tf.shape(x)[1]])
        for label in range(10):
            h, label = self.overlay_y_on_x(data=(x, label))
            h = tf.reshape(h, [-1, tf.shape(h)[0]])
            goodness = []
            for layer_idx in range(1, len(self.layer_list)):
                layer = self.layer_list[layer_idx]
                h = layer(h)
                goodness += [tf.math.reduce_mean(tf.math.pow(h, 2), 1)]
            goodness_per_label += [
                tf.expand_dims(tf.reduce_sum(goodness, keepdims=True), 1)
            ]
        goodness_per_label = tf.concat(goodness_per_label, 1)
        return tf.to_float(tf.argmax(goodness_per_label, 1))
//  [python]
def predict_one_sample(self, x):
        goodness_per_label = []
        x = tf.reshape(x, [tf.shape(x)[0] * tf.shape(x)[1]])
        for label in range(10):
            h, label = self.overlay_y_on_x(data=(x, label))
            h = tf.reshape(h, [-1, tf.shape(h)[0]])
            goodness = []
            for layer_idx in range(1, len(self.layer_list)):
                layer = self.layer_list[layer_idx]
                h = layer(h)
                goodness += [tf.math.reduce_mean(tf.math.pow(h, 2), 1)]
            goodness_per_label += [
                tf.expand_dims(tf.reduce_sum(goodness, keepdims=True), 1)
            ]
        goodness_per_label = tf.concat(goodness_per_label, 1)
        return tf.cast(tf.argmax(goodness_per_label, 1), tf.float64)
//  [python]
def test_output_equivalence(self):
        for tokenizer, tf_tokenizer in zip(self.tokenizers, self.tf_tokenizers):
            for test_inputs in self.test_sentences:
                python_outputs = tokenizer([test_inputs], return_tensors="tf")
                tf_outputs = tf_tokenizer([test_inputs])

                for key in python_outputs.keys():
                    # convert them to numpy to avoid messing with ragged tensors
                    python_outputs_values = python_outputs[key].numpy()
                    tf_outputs_values = tf_outputs[key].numpy()

                    self.assertTrue(tf.reduce_all(python_outputs_values.shape == tf_outputs_values.shape))
                    self.assertTrue(tf.reduce_all(tf.cast(python_outputs_values, tf.int64) == tf_outputs_values))
//  [python]
def decode_record(record, name_to_features):
  """Decodes a record to a TensorFlow example."""
  example = tf.io.parse_single_example(record, name_to_features)

  # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
  # So cast all int64 to int32.
  for name in list(example.keys()):
    t = example[name]
    if t.dtype == tf.int64:
      t = t.astype(tf.int32)
    example[name] = t

  return example
//  [python]
def decode_record(record, name_to_features):
  """Decodes a record to a TensorFlow example."""
  example = tf.io.parse_single_example(record, name_to_features)

  # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
  # So cast all int64 to int32.
  for name in list(example.keys()):
    t = example[name]
    if t.dtype == tf.int64:
      t = tf.cast(t, tf.int32)
    example[name] = t

  return example