tfds.features.Audio eagerly imports TensorFlow in _AudioDecoder.__init__, breaking TensorFlow-less usage

[owahltinez]

Description

TFDS supports TensorFlow-less dataset reading (e.g. via tfds.data_source with JAX, PyTorch, or Grain). However, constructing any dataset containing an audio feature (such as librispeech) or directly instantiating tfds.features.Audio fails with ModuleNotFoundError: No module named 'tensorflow' in an environment without TensorFlow.

Root Cause

In tensorflow_datasets/core/features/audio_feature.py, _AudioDecoder.__init__ eagerly evaluates tf.dtypes.as_dtype:

class _AudioDecoder(abc.ABC):
  def __init__(self, file_format: Optional[str], np_dtype: np.dtype, shape: utils.Shape):
    self._file_format = file_format
    self._np_dtype = np_dtype
    self._dtype = tf.dtypes.as_dtype(self._np_dtype)  # <--- Forces import of lazy tf module
    self._shape = shape
    self._channels = shape[1] if len(shape) > 1 else 1

Because tf is lazily imported via etils.epy.lazy_imports, accessing tf.dtypes triggers import tensorflow immediately during feature specification construction in _info(), even when only reading NumPy arrays.

Minimal Reproduction

In a virtual environment with only tensorflow-datasets installed (no tensorflow):

import tensorflow_datasets as tfds

# Fails with ModuleNotFoundError: No module named 'tensorflow'
tfds.features.Audio(sample_rate=16000)
# Or:
tfds.builder("librispeech")

Minimal Workaround

Monkeypatch _AudioDecoder.__init__ before importing tfds:

import tensorflow_datasets.core.features.audio_feature as af

def _patched_init(self, file_format, np_dtype, shape):
  self._file_format = file_format
  self._np_dtype = np_dtype
  self._dtype = None
  self._shape = shape
  self._channels = shape[1] if len(shape) > 1 else 1

af._AudioDecoder.__init__ = _patched_init

import tensorflow_datasets as tfds
# Now succeeds without TensorFlow installed:
ds = tfds.data_source("librispeech", split="train_clean100")

Suggested Fix

Make _dtype a lazy property or defer tf.dtypes.as_dtype to _LazyDecoder.decode_audio, where TensorFlow tensors are actually consumed:

@property
def _dtype(self):
  return tf.dtypes.as_dtype(self._np_dtype)