Add batch support for resample

tensorflow_io/core/python/ops/audio_ops.py

@@ -372,12 +372,13 @@ def fade(input, fade_in, fade_out, mode, name=None):
     return factor_in * factor_out * input
 
 
-def resample(input, rate_in, rate_out, name=None):  # pylint: disable=redefined-builtin
+def resample(input, rate_in, rate_out, name=None):
     """Resample audio.
 
     Args:
-      input: A 1-D (`[samples]`) or 2-D (`[samples, channels]`) `Tensor` of
-        type `int16` or `float`. Audio input.
+      input: A 1-D (`[samples]`) or 2-D (`[samples, channels]`) or 3-D
+        (`[batch, samples, channels]`) `Tensor` of type
+        `int16` or `float`. Audio input.
       rate_in: The rate of the audio input.
       rate_out: The rate of the audio output.
       name: A name for the operation (optional).
@@ -387,14 +388,37 @@ def resample(input, rate_in, rate_out, name=None):  # pylint: disable=redefined-
     """
     rank = tf.rank(input)
 
-    input = tf.cond(
-        tf.math.equal(rank, 1), lambda: tf.expand_dims(input, -1), lambda: input
-    )
-    value = core_ops.io_audio_resample(
-        input, rate_in=rate_in, rate_out=rate_out, name=name
+    def f1():
+        return tf.expand_dims(tf.expand_dims(input, -1), 0)
+
+    def f2():
+        return tf.expand_dims(input, 0)
+
+    def f3():
+        return input
+
+    input = tf.case(
+        [(tf.math.equal(rank, 1), f1), (tf.math.equal(rank, 2), f2)], default=f3
     )
-    return tf.cond(
-        tf.math.equal(rank, 1), lambda: tf.squeeze(value, [-1]), lambda: value
+
+    def f(i):
+        return core_ops.io_audio_resample(
+            i, rate_in=rate_in, rate_out=rate_out, name=name
+        )
+
+    value = tf.vectorized_map(f, input)
+
+    def g1():
+        return tf.squeeze(value, [0, -1])
+
+    def g2():
+        return tf.squeeze(value, [0])
+
+    def g3():
+        return value
+
+    return tf.case(
+        [(tf.math.equal(rank, 1), g1), (tf.math.equal(rank, 2), g2)], default=g3
     )
 
 

tests/test_audio_ops.py

@@ -72,6 +72,13 @@ def fixture_resample_1d():
     return args[:, 0], func, expected[:, 0]
 
 
+@pytest.fixture(name="resample_batch", scope="module")
+def fixture_resample_batch():
+    """fixture_resample_batch"""
+    args, func, expected = fixture_resample_base()
+    return tf.stack([args, args]), func, tf.stack([expected, expected])
+
+
 @pytest.fixture(name="decode_wav", scope="module")
 def fixture_decode_wav():
     """fixture_decode_wav"""
@@ -633,6 +640,7 @@ def func(e):
     [
         pytest.param("resample"),
         pytest.param("resample_1d"),
+        pytest.param("resample_batch"),
         pytest.param("decode_wav"),
         pytest.param("encode_wav"),
         pytest.param("decode_wav_u8"),
@@ -683,6 +691,7 @@ def func(e):
     ids=[
         "resample",
         "resample[1d]",
+        "resample[batch]",
         "decode_wav",
         "encode_wav",
         "decode_wav|u8",