Fix torchcodec audio decoding to respect 'num_channels'

src/datasets/features/_torchcodec.py

@@ -6,7 +6,14 @@ class AudioDecoder(_AudioDecoder):
     def __getitem__(self, key: str):
         if key == "array":
             y = self.get_all_samples().data.cpu().numpy()
-            return np.mean(y, axis=tuple(range(y.ndim - 1))) if y.ndim > 1 else y
+            if y.ndim <= 1:
+                return y
+            requested_num_channels = getattr(self, "_hf_num_channels", None)
+            # Backward compatibility: default behavior returns mono unless
+            # num_channels was explicitly set to a multi-channel value.
+            if requested_num_channels is None or requested_num_channels == 1:
+                return np.mean(y, axis=tuple(range(y.ndim - 1)))
+            return y
         elif key == "sampling_rate":
             return self.get_samples_played_in_range(0, 0).sample_rate
         elif hasattr(super(), "__getitem__"):

src/datasets/features/audio.py

@@ -216,6 +216,7 @@ def decode_example(
             audio = AudioDecoder(
                 bytes, stream_index=self.stream_index, sample_rate=self.sampling_rate, num_channels=self.num_channels
             )
+        audio._hf_num_channels = self.num_channels
         audio._hf_encoded = {"path": path, "bytes": bytes}
         audio.metadata.path = path
         return audio

tests/features/test_audio.py

@@ -801,6 +801,8 @@ def test_audio_decode_example_opus_convert_to_stereo(shared_datadir):
     decoded_example = audio.decode_example(audio.encode_example(audio_path))
     assert isinstance(decoded_example, AudioDecoder)
     samples = decoded_example.get_all_samples()
+    assert decoded_example["array"].ndim == 2
+    assert decoded_example["array"].shape[0] == 2
     assert samples.sample_rate == 48000
     assert samples.data.shape == (2, 48000)
 
@@ -815,5 +817,7 @@ def test_audio_decode_example_opus_convert_to_mono(shared_datadir):
     decoded_example = audio.decode_example(audio.encode_example(audio_path))
     assert isinstance(decoded_example, AudioDecoder)
     samples = decoded_example.get_all_samples()
+    assert decoded_example["array"].ndim == 1
+    assert abs(decoded_example["array"].shape[0] - samples.data.shape[1]) < 2
     assert samples.sample_rate == 44100
     assert samples.data.shape == (1, 202311)