Don't import TensorFlow when defining an Audio feature

tensorflow_datasets/core/features/audio_feature.py

@@ -82,10 +82,14 @@ def __init__(
     """
     self._file_format = file_format
     self._np_dtype = np_dtype
-    self._dtype = tf.dtypes.as_dtype(self._np_dtype)
     self._shape = shape
     self._channels = shape[1] if len(shape) > 1 else 1
 
+  @property
+  def _dtype(self) -> tf.dtypes.DType:
+    # Computed on access so that defining the feature doesn't import TensorFlow.
+    return tf.dtypes.as_dtype(self._np_dtype)
+
   @abc.abstractmethod
   def encode_audio(
       self, fobj: BinaryIO, file_format: Optional[str]

tensorflow_datasets/core/features/audio_feature_test.py

@@ -18,13 +18,15 @@
 import array
 import pathlib
 import tempfile
+from unittest import mock
 
 from absl.testing import parameterized
 import numpy as np
 import pydub
 import tensorflow as tf
 from tensorflow_datasets import testing
 from tensorflow_datasets.core import features
+from tensorflow_datasets.core.features import audio_feature
 
 
 class AudioFeatureTest(
@@ -190,6 +192,14 @@ def read(self, *args, **kwargs):
           ],
       )
 
+  def test_construction_does_not_use_tensorflow(self):
+    # Defining the feature must not import TensorFlow, so TFDS can be used
+    # without it (e.g. with `tfds.data_source`).
+    with mock.patch.object(audio_feature, 'tf') as tf_mock:
+      features.Audio(sample_rate=1000)
+      features.Audio(sample_rate=1000, file_format='wav', shape=(None, 2))
+    tf_mock.dtypes.as_dtype.assert_not_called()
+
   def test_sample_rate_property(self):
     self.assertEqual(features.Audio(sample_rate=1600).sample_rate, 1600)