def _num_tokens(self, data):
        """Count the number of tokens in a ragged, sparse or dense tensor."""
        if isinstance(data, tf.SparseTensor):
            flat_values = data.values
        elif isinstance(data, tf.RaggedTensor):
            flat_values = data.flat_values
        else:
            flat_values = tf.reshape(data, [-1])
        tokens, _, counts = tf.unique_with_counts(flat_values, out_idx="int64")
        return tokens, counts
Count the number of tokens in a ragged, sparse or dense tensor.
keras/src/layers/preprocessing/index_lookup.py