Adds FFT Auto Scale Op (#2134)

This PR adds additional FFT op functionality in the Signal library, namely adding the FFT Auto Scale operation.
Testing added in the original `fft_test.cc` and `fft_ops_test.py`.
BUG=[287346710](http://b/287346710)

python/tflite_micro/python_ops_resolver.cc

@@ -51,6 +51,7 @@ PythonOpsResolver::PythonOpsResolver() {
  AddEthosU();
  AddExp();
  AddExpandDims();
+ AddFftAutoScale();
  AddFill();
  AddFloor();
  AddFloorDiv();

python/tflite_micro/signal/BUILD

@@ -46,6 +46,10 @@ py_tflm_signal_library(
 py_test(
  name = "fft_ops_test",
  srcs = ["ops/fft_ops_test.py"],
+ data = [
+ "//python/tflite_micro/signal/ops/testdata:fft_auto_scale_test1.txt",
+ "//python/tflite_micro/signal/ops/testdata:rfft_test1.txt",
+ ],
  python_version = "PY3",
  srcs_version = "PY3",
  deps = [

python/tflite_micro/signal/ops/fft_ops.py

@@ -65,5 +65,22 @@ def _fft(input_tensor, fft_length, name=default_name):
  return _fft
 
 
+def _fft_auto_scale_wrapper(fft_auto_scale_fn, default_name):
+ """Wrapper around gen_fft_ops.fft_auto_scale*."""
+
+ def _fft_auto_scale(input_tensor, name=default_name):
+ with tf.name_scope(name) as name:
+ input_tensor = tf.convert_to_tensor(input_tensor, dtype=tf.int16)
+ dim_list = input_tensor.shape.as_list()
+ if len(dim_list) != 1:
+ raise ValueError("Input tensor must have a rank of 1")
+ return fft_auto_scale_fn(input_tensor, name=name)
+
+ return _fft_auto_scale
+
+
 rfft = _fft_wrapper(gen_fft_ops.signal_rfft, "signal_rfft")
+fft_auto_scale = _fft_auto_scale_wrapper(gen_fft_ops.signal_fft_auto_scale,
+ "signal_fft_auto_scale")
 tf.no_gradient("signal_rfft")
+tf.no_gradient("signal_fft_auto_scale")

python/tflite_micro/signal/ops/fft_ops_test.py

@@ -33,6 +33,31 @@ def GetResource(self, filepath):
  file_text = f.read()
  return file_text
 
+ def SingleFftAutoScaleTest(self, filename):
+ lines = self.GetResource(filename).splitlines()
+ func = tf.function(fft_ops.fft_auto_scale)
+ input_size = len(lines[0].split())
+ concrete_function = func.get_concrete_function(
+ tf.TensorSpec(input_size, dtype=tf.int16))
+ interpreter = util.get_tflm_interpreter(concrete_function, func)
+ i = 0
+ while i < len(lines):
+ in_frame = np.array([int(j) for j in lines[i].split()], dtype=np.int16)
+ out_frame_exp = [int(j) for j in lines[i + 1].split()]
+ scale_exp = [int(j) for j in lines[i + 2].split()]
+ # TFLM
+ interpreter.set_input(in_frame, 0)
+ interpreter.invoke()
+ out_frame = interpreter.get_output(0)
+ scale = interpreter.get_output(1)
+ self.assertAllEqual(out_frame_exp, out_frame)
+ self.assertEqual(scale_exp, scale)
+ # TF
+ out_frame, scale = self.evaluate(fft_ops.fft_auto_scale(in_frame))
+ self.assertAllEqual(out_frame_exp, out_frame)
+ self.assertEqual(scale_exp, scale)
+ i += 3
+
  def SingleRfftTest(self, filename):
  lines = self.GetResource(filename).splitlines()
  args = lines[0].split()
@@ -43,8 +68,6 @@ def SingleRfftTest(self, filename):
  tf.TensorSpec(input_size, dtype=tf.int16), fft_length)
  # TODO(b/286252893): make test more robust (vs scipy)
  interpreter = util.get_tflm_interpreter(concrete_function, func)
- input_details = interpreter.get_input_details()
- output_details = interpreter.get_output_details()
  # Skip line 0, which contains the configuration params.
  # Read lines in pairs <input, expected>
  i = 1
@@ -53,9 +76,9 @@ def SingleRfftTest(self, filename):
  out_frame_exp = [int(j) for j in lines[i + 1].split()]
  # Compare TFLM inference against the expected golden values
  # TODO(b/286252893): validate usage of testing vs interpreter here
- interpreter.set_tensor(input_details[0]['index'], in_frame)
+ interpreter.set_input(in_frame, 0)
  interpreter.invoke()
- out_frame = interpreter.get_tensor(output_details[0]['index'])
+ out_frame = interpreter.get_output(0)
  self.assertAllEqual(out_frame_exp, out_frame)
  # TF
  out_frame = self.evaluate(fft_ops.rfft(in_frame, fft_length))
@@ -83,11 +106,9 @@ def MultiDimRfftTest(self, filename):
  concrete_function = func.get_concrete_function(
  tf.TensorSpec(np.shape(in_frames), dtype=tf.int16), fft_length)
  interpreter = util.get_tflm_interpreter(concrete_function, func)
- input_details = interpreter.get_input_details()
- output_details = interpreter.get_output_details()
- interpreter.set_tensor(input_details[0]['index'], in_frames)
+ interpreter.set_input(in_frames, 0)
  interpreter.invoke()
- out_frame = interpreter.get_tensor(output_details[0]['index'])
+ out_frame = interpreter.get_output(0)
  self.assertAllEqual(out_frames_exp, out_frame)
  # TF
  out_frames = self.evaluate(fft_ops.rfft(in_frames, fft_length))
@@ -204,6 +225,12 @@ def testRfftSineTest(self):
  delta=1)
  fft_length = 2 * fft_length
 
+ def testRfft(self):
+ self.SingleRfftTest('testdata/rfft_test1.txt')
+
+ def testRfftLargeOuterDimension(self):
+ self.MultiDimRfftTest('testdata/rfft_test1.txt')
+
  def testFftTooLarge(self):
  for dtype in [np.int16, np.int32, np.float32]:
  fft_input = np.zeros(round(fft_ops._MAX_FFT_LENGTH * 2), dtype=dtype)
@@ -224,6 +251,9 @@ def testFftLengthNoEven(self):
  with self.assertRaises((tf.errors.InvalidArgumentError, ValueError)):
  self.evaluate(fft_ops.rfft(fft_input, 127))
 
+ def testAutoScale(self):
+ self.SingleFftAutoScaleTest('testdata/fft_auto_scale_test1.txt')
+
  def testPow2FftLengthTest(self):
  fft_length, fft_bits = fft_ops.get_pow2_fft_length(131)
  self.assertEqual(fft_length, 256)

python/tflite_micro/signal/ops/testdata/BUILD

@@ -7,6 +7,7 @@ package(
 )
 
 exports_files([
+ "fft_auto_scale_test1.txt",
  "rfft_test1.txt",
  "window_test1.txt",
 ])