Repository: deepjavalibrary/djl
Pull request #2259 (2022-12-23 22:12:16)

[Audio] add stft and fft forier transform
## Description ##

Build forier transform operator and complex64 type support on DJL.


Patch:
diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index e82ddc80b..502ccf535 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -1,6 +1,9 @@
 ARG VARIANT=11
 FROM mcr.microsoft.com/vscode/devcontainers/java:${VARIANT}
 
+# Allow to use the VARIANT after FROM
+ARG VARIANT
+
 # [Optional] Install Maven or Gradle
 ARG INSTALL_MAVEN="false"
 ARG MAVEN_VERSION=3.6.3
@@ -14,6 +17,11 @@ ARG INSTALL_NODE="true"
 ARG NODE_VERSION="lts/*"
 RUN if [ "${INSTALL_NODE}" = "true" ]; then su vscode -c "source /usr/local/share/nvm/nvm.sh && nvm install ${NODE_VERSION} 2>&1"; fi
 
+# Install the lost package `fontconfig` for VARIANT 11
+RUN if [ "${VARIANT}" = "11" ]; then \
+    apt-get update && export DEBIAN_FRONTEND=noninteractive \
+    && apt-get -y install --no-install-recommends fontconfig; fi
+
 # [Optional] Uncomment this section to install additional OS packages.
 # RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
 #     && apt-get -y install --no-install-recommends <your-package-list-here>
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 59ec56534..ee67c7f78 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -3195,6 +3195,65 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray flatten(int startDim, int endDim);
 
+    /**
+     * Computes the one-dimensional discrete Fourier Transform.
+     *
+     * @param length Length of the transformed axis of the output.
+     * @return The truncated or zero-padded input, transformed along the axis indicated by axis, or
+     *     the last one if axis is not specified.
+     */
+    default NDArray fft(long length) {
+        return fft(length, -1);
+    }
+
+    /**
+     * Computes the one-dimensional discrete Fourier Transform.
+     *
+     * @param length Length of the transformed axis of the output.
+     * @param axis Axis over which to compute the FFT.
+     * @return The truncated or zero-padded input, transformed along the axis indicated by axis, or
+     *     the last one if axis is not specified.
+     */
+    NDArray fft(long length, long axis);
+
+    /**
+     * Computes the Short Time Fourier Transform (STFT).
+     *
+     * @param nFft size of Fourier transform
+     * @param hopLength the distance between neighboring sliding window frames. Default can be:
+     *     floor(n_fft / 4)
+     * @param center whether to pad input on both sides.
+     * @param window Desired window to use. Recommend for HanningWindow
+     * @param returnComplex whether to return a complex tensor, or a real tensor with an extra last
+     *     dimension for the real and imaginary components.
+     * @return A NDArray containing the STFT result with shape described above
+     */
+    default NDArray stft(
+            long nFft, long hopLength, boolean center, NDArray window, boolean returnComplex) {
+        return stft(nFft, hopLength, center, window, false, returnComplex);
+    }
+
+    /**
+     * Computes the Short Time Fourier Transform (STFT).
+     *
+     * @param nFft size of Fourier transform
+     * @param hopLength the distance between neighboring sliding window frames. Default can be:
+     *     floor(n_fft / 4)
+     * @param center whether to pad input on both sides.
+     * @param window Desired window to use. Recommend for HanningWindow
+     * @param normalize controls whether to return the normalized STFT results
+     * @param returnComplex whether to return a complex tensor, or a real tensor with an extra last
+     *     dimension for the real and imaginary components.
+     * @return A NDArray containing the STFT result with shape described above
+     */
+    NDArray stft(
+            long nFft,
+            long hopLength,
+            boolean center,
+            NDArray window,
+            boolean normalize,
+            boolean returnComplex);
+
     /**
      * Reshapes this {@code NDArray} to the given {@link Shape}.
      *
@@ -4978,4 +5037,21 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @return the result {@code NDArray}
      */
     NDArray batchDot(NDArray other);
+
+    /**
+     * Convert a general NDArray to its complex math format.
+     *
+     * <p>example: [10f, 12f] float32 -&gt; [10+12j] in complex64
+     *
+     * @return the complex NDArray
+     */
+    NDArray complex();
+
+    /**
+     * Convert a complex NDArray to its real math format. example: [10+12j] in complex64 -&gt; [10f,
+     * 12f] float32
+     *
+     * @return tje real NDArray
+     */
+    NDArray real();
 }
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index fbe2d7b05..5ed86277c 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -834,6 +834,24 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray stft(
+            long nFft,
+            long hopLength,
+            boolean center,
+            NDArray window,
+            boolean normalize,
+            boolean returnComplex) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
@@ -1122,6 +1140,18 @@ public abstract class NDArrayAdapter implements NDArray {
         return getAlternativeArray().batchDot(other);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray complex() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray real() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArrayEx getNDArrayInternal() {
diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index 26ffcd7aa..608dfe85f 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -1461,6 +1461,25 @@ public interface NDManager extends AutoCloseable {
      */
     NDArray sampleGamma(NDArray alpha, NDArray beta, Shape shape);
 
+    /**
+     * Builds the Hanning Window.
+     *
+     * <p>The Hanning was named for Julius von Hann, an Austrian meteorologist. It is also known as
+     * the Cosine Bell. Some authors prefer that it be called a Hann window, to help avoid confusion
+     * with the very similar Hamming window.
+     *
+     * @param numPoints Number of points in the output window.
+     * @return the window
+     */
+    default NDArray hanningWindow(long numPoints) {
+        float[] data = new float[(int) numPoints];
+        // shift from N -1 to N to trims off the last duplicate value from the symmetric window
+        for (int i = 1; i < data.length; i++) {
+            data[i] = (float) (0.5 * (1 - Math.cos((2 * Math.PI * i) / numPoints)));
+        }
+        return create(data);
+    }
+
     /**
      * Check if the manager is still valid.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/types/DataType.java b/api/src/main/java/ai/djl/ndarray/types/DataType.java
index 620cad66f..e4af63081 100644
--- a/api/src/main/java/ai/djl/ndarray/types/DataType.java
+++ b/api/src/main/java/ai/djl/ndarray/types/DataType.java
@@ -33,6 +33,7 @@ public enum DataType {
     INT8(Format.INT, 1),
     INT64(Format.INT, 8),
     BOOLEAN(Format.BOOLEAN, 1),
+    COMPLEX64(Format.FLOATING, 4),
     UNKNOWN(Format.UNKNOWN, 0),
     STRING(Format.STRING, -1);
 
@@ -185,6 +186,7 @@ public enum DataType {
                 return data.asLongBuffer();
             case UINT8:
             case INT8:
+            case COMPLEX64:
             case UNKNOWN:
             default:
                 return data;
@@ -217,6 +219,7 @@ public enum DataType {
                 return "|b1";
             case STRING:
                 return "|S1";
+            case COMPLEX64:
             case UNKNOWN:
             default:
                 throw new IllegalArgumentException("Unsupported dataType: " + this);
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index 8f9dcbc0b..479a1a678 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -1104,6 +1104,24 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray stft(
+            long nFft,
+            long hopLength,
+            boolean center,
+            NDArray window,
+            boolean normalize,
+            boolean returnComplex) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
@@ -1556,6 +1574,18 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         return manager.invoke("_npx_batch_dot", new NDArray[] {this, other}, null);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray complex() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray real() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArrayEx getNDArrayInternal() {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index d51cc0acd..de96b24f2 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -1056,6 +1056,25 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         return JniUtils.flatten(this, startDim, endDim);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft(long length, long axis) {
+        return JniUtils.fft(this, length, axis);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray stft(
+            long nFft,
+            long hopLength,
+            boolean center,
+            NDArray window,
+            boolean normalize,
+            boolean returnComplex) {
+        return JniUtils.stft(
+                this, nFft, hopLength, (PtNDArray) window, center, normalize, returnComplex);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDArray reshape(Shape shape) {
@@ -1512,6 +1531,18 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         throw new UnsupportedOperationException("Not implemented");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray complex() {
+        return JniUtils.complex(this);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray real() {
+        return JniUtils.real(this);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDArrayEx getNDArrayInternal() {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
index 71f0cb57f..0188acea3 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
@@ -172,6 +172,12 @@ public class PtNDManager extends BaseNDManager {
         return JniUtils.normal(this, loc, scale, shape, dataType, device);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray hanningWindow(long numPoints) {
+        return JniUtils.hannWindow(this, numPoints, true, device);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDManager newSubManager(Device device) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 411c4843a..741d12fc7 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -924,6 +924,42 @@ public final class JniUtils {
                 PyTorchLibrary.LIB.torchFlatten(ndArray.getHandle(), startDim, endDim));
     }
 
+    public static PtNDArray fft(PtNDArray ndArray, long length, long axis) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchFft(ndArray.getHandle(), length, axis));
+    }
+
+    public static PtNDArray stft(
+            PtNDArray ndArray,
+            long nFft,
+            long hopLength,
+            PtNDArray window,
+            boolean center,
+            boolean normalize,
+            boolean returnComplex) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchStft(
+                        ndArray.getHandle(),
+                        nFft,
+                        hopLength,
+                        window.getHandle(),
+                        center,
+                        normalize,
+                        returnComplex));
+    }
+
+    public static PtNDArray real(PtNDArray ndArray) {
+        return new PtNDArray(
+                ndArray.getManager(), PyTorchLibrary.LIB.torchViewAsReal(ndArray.getHandle()));
+    }
+
+    public static PtNDArray complex(PtNDArray ndArray) {
+        return new PtNDArray(
+                ndArray.getManager(), PyTorchLibrary.LIB.torchViewAsComplex(ndArray.getHandle()));
+    }
+
     public static PtNDArray abs(PtNDArray ndArray) {
         return new PtNDArray(
                 ndArray.getManager(), PyTorchLibrary.LIB.torchAbs(ndArray.getHandle()));
@@ -1187,6 +1223,16 @@ public final class JniUtils {
                         false));
     }
 
+    public static PtNDArray hannWindow(
+            PtNDManager manager, long numPoints, boolean periodic, Device device) {
+        return new PtNDArray(
+                manager,
+                PyTorchLibrary.LIB.torchHannWindow(
+                        numPoints,
+                        periodic,
+                        new int[] {PtDeviceType.toDeviceType(device), device.getDeviceId()}));
+    }
+
     public static PtNDArray erfinv(PtNDArray ndArray) {
         return new PtNDArray(
                 ndArray.getManager(), PyTorchLibrary.LIB.torchErfinv(ndArray.getHandle()));
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 7196b55be..f6c4b3d2c 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -248,6 +248,21 @@ final class PyTorchLibrary {
 
     native long torchFlatten(long handle, long startDim, long endDim);
 
+    native long torchFft(long handle, long length, long axis);
+
+    native long torchStft(
+            long handle,
+            long nFft,
+            long hopLength,
+            long windowHandle,
+            boolean center,
+            boolean normalize,
+            boolean returnComplex);
+
+    native long torchViewAsReal(long handle);
+
+    native long torchViewAsComplex(long handle);
+
     native long[] torchSplit(long handle, long size, long dim);
 
     native long[] torchSplit(long handle, long[] indices, long dim);
@@ -369,6 +384,8 @@ final class PyTorchLibrary {
 
     native long torchEye(int n, int m, int dType, int layout, int[] device, boolean requiredGrad);
 
+    native long torchHannWindow(long nfft, boolean periodic, int[] device);
+
     native long torchErfinv(long handle);
 
     native long torchInverse(long self);
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc
index cb86295e1..ebf015014 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc
@@ -159,3 +159,14 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchSparseCoo(
   return reinterpret_cast<uintptr_t>(result_ptr);
   API_END_RETURN();
 }
+
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchHannWindow(
+        JNIEnv* env, jobject jthis, jlong jn_fft, jboolean jperoidic, jintArray jdevice) {
+    API_BEGIN()
+        const auto device = utils::GetDeviceFromJDevice(env, jdevice);
+        auto options = torch::TensorOptions().device(device);
+        const auto* result_ptr = new torch::Tensor(torch::hann_window(jn_fft, jperoidic, options));
+        return reinterpret_cast<uintptr_t>(result_ptr);
+    API_END_RETURN()
+}
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
index f24fa0a91..8dda3cdc3 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
@@ -24,3 +24,46 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFlatten(
   return reinterpret_cast<uintptr_t>(result_ptr);
   API_END_RETURN()
 }
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFft(
+        JNIEnv* env, jobject jthis, jlong jhandle, jlong jn, jlong jaxis) {
+    API_BEGIN()
+        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+        const auto* result_ptr = new torch::Tensor(torch::fft_fft(*tensor_ptr, jn, jaxis));
+        return reinterpret_cast<uintptr_t>(result_ptr);
+    API_END_RETURN()
+}
+
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchStft(
+        JNIEnv* env, jobject jthis, jlong jhandle, jlong jn_fft, jlong jhop_length, jlong jwindow, jboolean jcenter, jboolean jnormalize, jboolean jreturn_complex) {
+    API_BEGIN()
+        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+        if (jwindow == -1L) {
+            const auto* result_ptr = new torch::Tensor(tensor_ptr->stft(jn_fft, jhop_length, c10::nullopt, c10::nullopt, jcenter, "reflect", jnormalize, c10::nullopt, jreturn_complex));
+            return reinterpret_cast<uintptr_t>(result_ptr);
+        } else {
+            const auto* window_ptr = reinterpret_cast<torch::Tensor*>(jwindow);
+            const auto* result_ptr = new torch::Tensor(tensor_ptr->stft(jn_fft, jhop_length, c10::nullopt, *window_ptr, jcenter, "reflect", jnormalize, c10::nullopt, jreturn_complex));
+            return reinterpret_cast<uintptr_t>(result_ptr);
+        }
+    API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchViewAsReal(
+        JNIEnv* env, jobject jthis, jlong jhandle) {
+    API_BEGIN()
+        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+        const auto* result_ptr = new torch::Tensor(torch::view_as_real_copy(*tensor_ptr));
+        return reinterpret_cast<uintptr_t>(result_ptr);
+    API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchViewAsComplex(
+        JNIEnv* env, jobject jthis, jlong jhandle) {
+    API_BEGIN()
+        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+        const auto* result_ptr = new torch::Tensor(torch::view_as_complex_copy(*tensor_ptr));
+        return reinterpret_cast<uintptr_t>(result_ptr);
+    API_END_RETURN()
+}
diff --git a/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h b/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
index 45fd527ff..dc3edb099 100644
--- a/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
+++ b/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
@@ -81,6 +81,8 @@ inline torch::ScalarType GetScalarTypeFromDType(jint dtype) {
       return torch::kInt64;
     case 7:
       return torch::kBool;
+    case 8:
+      return torch::kComplexFloat;
     default:
       // TODO improve the error handling
       throw;
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index 0cf130c1e..f19c754ac 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -524,6 +524,18 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray complex() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray real() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray add(Number n) {
@@ -1140,6 +1152,24 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray stft(
+            long nFft,
+            long hopLength,
+            boolean center,
+            NDArray window,
+            boolean normalize,
+            boolean returnComplex) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/LinearSpecgram.java b/extensions/audio/src/main/java/ai/djl/audio/processor/LinearSpecgram.java
index 7f1e7ae79..e708788d0 100644
--- a/extensions/audio/src/main/java/ai/djl/audio/processor/LinearSpecgram.java
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/LinearSpecgram.java
@@ -63,7 +63,7 @@ public class LinearSpecgram implements AudioProcessor {
         }
         samples = NDArrays.stack(windowList);
 
-        NDArray weighting = hanningWindow(windowSize, manager);
+        NDArray weighting = manager.hanningWindow(windowSize);
         samples.muli(weighting);
         NDList fftList = new NDList();
         for (int row = 0; row < rows; row++) {
@@ -140,12 +140,4 @@ public class LinearSpecgram implements AudioProcessor {
         }
         return in.getManager().create(abs);
     }
-
-    private NDArray hanningWindow(int size, NDManager manager) {
-        float[] data = new float[size];
-        for (int i = 1; i < size; i++) {
-            data[i] = (float) (0.5 * (1 - Math.cos((2 * Math.PI * i) / (size - 1))));
-        }
-        return manager.create(data);
-    }
 }


Test patch:
diff --git a/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
index 02c31aab5..8070a1842 100644
--- a/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
+++ b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
@@ -64,6 +64,6 @@ public class AudioProcessorTest {
         testData.setAudioPaths(Collections.singletonList("build/test/test_01.wav"));
         NDArray samples = testData.getPreprocessedData(manager, 0);
         Assert.assertEquals(new Shape(161, 838), samples.getShape());
-        Assert.assertEquals(samples.get("0,0").toFloatArray()[0], -15.4571f, 1e-3f);
+        Assert.assertEquals(samples.get("0,0").toFloatArray()[0], -15.4531f, 1e-3f);
     }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java
index 00b346eab..853142cb5 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayCreationOpTest.java
@@ -525,4 +525,13 @@ public class NDArrayCreationOpTest {
             Assertions.assertAlmostEquals(std.toFloatArray()[1], 13.69f, 2e-2f, 2e-2f);
         }
     }
+
+    @Test
+    public void testHanningWindow() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray result = manager.hanningWindow(6);
+            NDArray expected = manager.create(new float[] {0.0f, 0.25f, 0.75f, 1.0f, 0.75f, 0.25f});
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
index 3c7868b6b..814ec62f7 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
@@ -984,4 +984,66 @@ public class NDArrayOtherOpTest {
             }
         }
     }
+
+    @Test
+    public void testFft() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray ele = manager.create(new float[] {1, 2, 3, 4, 5});
+            NDArray result = ele.fft(5);
+            result = result.real().flatten();
+            NDArray expected =
+                    manager.create(
+                            new float[] {
+                                15.0f,
+                                0.0f,
+                                -2.5f,
+                                3.440955f,
+                                -2.4999998f,
+                                0.8122992f,
+                                -2.4999998f,
+                                -0.8122992f,
+                                -2.5f,
+                                -3.440955f
+                            });
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
+
+    @Test
+    public void testStft() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray audio = manager.ones(new Shape(20));
+            NDArray window = manager.hanningWindow(20);
+            int nfft = 20;
+            int hopLength = 16;
+            NDArray result = audio.stft(nfft, hopLength, false, window, true);
+            NDArray expected =
+                    manager.create(
+                            new float[] {
+                                10.0f,
+                                0.0f,
+                                -5.0000005f,
+                                -4.7683716E-7f,
+                                1.6598708E-8f,
+                                -6.70266E-8f,
+                                1.830015E-7f,
+                                -3.239836E-7f,
+                                4.651681E-7f,
+                                7.3675395E-8f,
+                                0.0f,
+                                -0.0f,
+                                1.166903E-8f,
+                                7.3675395E-8f,
+                                1.65798E-7f,
+                                -8.896347E-8f,
+                                -1.35808E-7f,
+                                1.5145238E-7f,
+                                3.5762787E-7f,
+                                0.0f,
+                                0.0f,
+                                0.0f
+                            });
+            Assertions.assertAlmostEquals(result.real().flatten(), expected);
+        }
+    }
 }