class TFFormatter(TensorFormatter[Mapping, "tf.Tensor", Mapping]):
    def __init__(self, features=None, token_per_repo_id=None, **tf_tensor_kwargs):
        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
        self.tf_tensor_kwargs = tf_tensor_kwargs
        import tensorflow as tf  # noqa: F401 - import tf at initialization

    def _consolidate(self, column):
        import tensorflow as tf

        if isinstance(column, list) and column:
            if all(
                isinstance(x, tf.Tensor) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column
            ):
                return tf.stack(column)
            elif all(
                isinstance(x, (tf.Tensor, tf.RaggedTensor)) and x.ndim == 1 and x.dtype == column[0].dtype
                for x in column
            ):
                # only rag 1-D tensors, otherwise some dimensions become ragged even though they were consolidated
                return tf.ragged.stack(column)