def select_from_act(self, act: tf.Tensor, label: tf.Tensor) -> tf.Tensor:
        """Takes tensor of acts act_{b, a, t, u, ...} and labels label_{b,u},
        where b is the batch index, t is the logit index, and u is the label index,
        and returns for each token index k the tensor

            output_{b,a,t,k,...} = logsumexp_u act_{b,a,t,u_k,...} * kroneker_delta(u_k = label_{b,u})

        that is logarithmic sum of exponents of acts for all u_k = label_{b,u}, given b, t and k.

        Args:
            act:    tf.Tensor, shape = [batch_size, dim_a, max_logit_length, max_label_length + 1, ...]
            label:  tf.Tensor, shape = [batch_size, max_label_length + 1]

        Returns:    tf.Tensor, shape = [batch_size, max_label_length + 1, num_tokens, ...]
        """
        data = smart_transpose(a=act, perm=[0, 3, 2, 1])
        # shape = [batch_size, max_label_length + 1, max_logit_length, dim_a, ...]
        data = tf.squeeze(
            input=smart_reshape(tensor=data, shape=[1, self.batch_size * self.max_label_length_plus_one, self.max_logit_length]), axis=0
        )
        # shape = [batch_size * (max_label_length + 1), max_logit_length, dim_a, ...]

        segment_ids = tf.reshape(label + tf.expand_dims(tf.range(self.batch_size), 1) * self.num_tokens, shape=[-1])
        # shape = [batch_size * (max_label_length + 1)]
        num_segments = self.batch_size * self.num_tokens

        output = unsorted_segment_logsumexp(data=data, segment_ids=segment_ids, num_segments=num_segments)
        # shape = [batch_size * num_tokens, max_logit_length, dim_a, ...]
        output = smart_reshape(tf.expand_dims(output, 0), [self.batch_size, self.num_tokens, self.max_logit_length])
        # shape = [batch_size, num_tokens, max_logit_length, dim_a, ...]
        output = smart_transpose(output, [0, 3, 2, 1])
        # shape = [batch_size, dim_a, max_logit_length, num_tokens, ...]
        return output
Takes tensor of acts act_{b, a, t, u, ...} and labels label_{b,u},
        where b is the batch index, t is the logit index, and u is the label index,
        and returns for each token index k the tensor

            output_{b,a,t,k,...} = logsumexp_u act_{b,a,t,u_k,...} * kroneker_delta(u_k = label_{b,u})

        that is logarithmic sum of exponents of acts for all u_k = label_{b,u}, given b, t and k.

        Args:
            act:    tf.Tensor, shape = [batch_size, dim_a, max_logit_length, max_label_length + 1, ...]
            label:  tf.Tensor, shape = [batch_size, max_label_length + 1]

        Returns:    tf.Tensor, shape = [batch_size, max_label_length + 1, num_tokens, ...]
        
tensorflow_asr/losses/impl/ctc_tpu.py