def METHOD_NAME(self, dataset_batch):
    """Returns the number of elements in a potentially partial batch."""
    if isinstance(dataset_batch, (tuple, list)):
        dataset_batch = dataset_batch[0]
    assert tf.nest.flatten(dataset_batch)
    def _find_any_tensor(batch_features):
        tensors = [
            x for x in tf.nest.flatten(batch_features) if tf.is_tensor(x)
        ]
        if not tensors:
            raise ValueError("Cannot find any Tensor in features dict.")
        return tensors[0]
    return backend.cast(
        backend.shape(_find_any_tensor(dataset_batch))[0], dtype="int64"
    )