Repository: deepjavalibrary/djl
Pull request #2845 (2023-11-09 16:19:18)

Add FFT2 and FFT2 inverse
## Description ##

Brief description of what this PR is about

- Added 2-dimensional FFT
- Added 2-dimensional inverse FFT


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1365 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.08% compared to head [(`5d3f508`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.29%.
> Report is 908 commits behind head on master.

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2845      +/-   ##
============================================
+ Coverage     72.08%   72.29%   +0.20%     
- Complexity     5126     7161    +2035     
============================================
  Files           473      708     +235     
  Lines         21970    31908    +9938     
  Branches       2351     3313     +962     
============================================
+ Hits          15838    23068    +7230     
- Misses         4925     7262    +2337     
- Partials       1207     1578     +371     
```


| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [...ava/ai/djl/inference/streaming/StreamingBlock.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL2luZmVyZW5jZS9zdHJlYW1pbmcvU3RyZWFtaW5nQmxvY2suamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Dimension.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9EaW1lbnNpb24uamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Unit.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9Vbml0LmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/audio/Audio.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2F1ZGlvL0F1ZGlvLmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...ava/ai/djl/modality/cv/output/DetectedObjects.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9EZXRlY3RlZE9iamVjdHMuamF2YQ==) | `96.29% <100.00%> (+1.29%)` | :arrow_up: |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <100.00%> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...i/djl/modality/cv/transform/RandomResizedCrop.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21SZXNpemVkQ3JvcC5qYXZh) | `94.11% <100.00%> (+5.22%)` | :arrow_up: |
| ... and [225 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [385 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2845?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 23d452be3..2b4a9df09 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -3393,6 +3393,48 @@ public interface NDArray extends NDResource, BytesSupplier {
             boolean normalize,
             boolean returnComplex);
 
+    /**
+     * Computes the two-dimensional Discrete Fourier Transform.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @param axes Axes over which to compute the 2D-FFT.
+     * @return The truncated or zero-padded input, transformed along the axes.
+     */
+    NDArray fft2(long[] sizes, long[] axes);
+
+    /**
+     * Computes the two-dimensional Discrete Fourier Transform along the last 2 axes.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @return The truncated or zero-padded input, transformed along the last two axes
+     */
+    default NDArray fft2(long[] sizes) {
+        return fft2(sizes, new long[] {-2, -1});
+    }
+
+    /**
+     * Computes the two-dimensional inverse Discrete Fourier Transform.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @param axes Axes over which to compute the 2D-Inverse-FFT.
+     * @return The truncated or zero-padded input, transformed along the axes.
+     */
+    NDArray ifft2(long[] sizes, long[] axes);
+
+    /**
+     * Computes the two-dimensional inverse Discrete Fourier Transform along the last 2 axes.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @return The truncated or zero-padded input, transformed along the axes.
+     */
+    default NDArray ifft2(long[] sizes) {
+        return ifft2(sizes, new long[] {-2, -1});
+    }
+
     /**
      * Reshapes this {@code NDArray} to the given {@link Shape}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 855c71830..9a4ad8db9 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -912,6 +912,18 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
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 19eec8372..8b884b399 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
@@ -1160,6 +1160,18 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
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 1e9ac83c1..606aaf24e 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
@@ -1103,6 +1103,18 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
                 this, nFft, hopLength, (PtNDArray) window, center, normalize, returnComplex);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        return JniUtils.fft2(this, sizes, axes);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        return JniUtils.ifft2(this, sizes, axes);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDArray reshape(Shape shape) {
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 8e6be7b8d..40a6a0065 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
@@ -1040,6 +1040,18 @@ public final class JniUtils {
         return new PtNDArray(ndArray.getManager(), handle);
     }
 
+    public static PtNDArray fft2(PtNDArray ndArray, long[] sizes, long[] axes) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchFft2(ndArray.getHandle(), sizes, axes));
+    }
+
+    public static PtNDArray ifft2(PtNDArray ndArray, long[] sizes, long[] axes) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchIfft2(ndArray.getHandle(), sizes, axes));
+    }
+
     public static PtNDArray real(PtNDArray ndArray) {
         long handle = PyTorchLibrary.LIB.torchViewAsReal(ndArray.getHandle());
         if (handle == -1) {
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 a1829306d..54fc54191 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
@@ -273,6 +273,10 @@ final class PyTorchLibrary {
             boolean normalize,
             boolean returnComplex);
 
+    native long torchFft2(long handle, long[] sizes, long[] axes);
+
+    native long torchIfft2(long handle, long[] sizes, long[] axes);
+
     native long torchViewAsReal(long handle);
 
     native long torchViewAsComplex(long handle);
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 5a65e1eca..08932098d 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
@@ -34,6 +34,28 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFft(
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFft2(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlongArray js, jlongArray jaxes) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const std::vector<int64_t> sizes = djl::utils::jni::GetVecFromJLongArray(env, js);
+  const std::vector<int64_t> axes = djl::utils::jni::GetVecFromJLongArray(env, jaxes);
+  const auto* result_ptr = new torch::Tensor(torch::fft_fft2(*tensor_ptr, sizes, axes));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIfft2(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlongArray js, jlongArray jaxes) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const std::vector<int64_t> sizes = djl::utils::jni::GetVecFromJLongArray(env, js);
+  const std::vector<int64_t> axes = djl::utils::jni::GetVecFromJLongArray(env, jaxes);
+  const auto* result_ptr = new torch::Tensor(torch::fft_ifft2(*tensor_ptr, sizes, axes));
+  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) {
 #ifdef V1_11_X
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 444885377..419be4c09 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
@@ -1184,6 +1184,18 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {


Test patch:
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 00e7465f7..66bb136ab 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
@@ -1087,4 +1087,58 @@ public class NDArrayOtherOpTest {
             Assertions.assertAlmostEquals(result.real().flatten(), expected);
         }
     }
+
+    @Test
+    public void testFft2() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray array =
+                    manager.create(
+                            new float[][] {
+                                {1f, 6.6f, 4.315f, 2.0f},
+                                {16.9f, 6.697f, 2.399f, 67.9f},
+                                {0f, 5f, 67.09f, 9.87f}
+                            });
+            NDArray result = array.fft2(new long[] {3, 4}, new long[] {0, 1});
+            result = result.real().flatten(1, 2); // flatten complex numbers
+            NDArray expected =
+                    manager.create(
+                            new float[][] {
+                                {189.771f, 0f, -55.904f, 61.473f, -6.363f, 0f, -55.904f, -61.473f},
+                                {
+                                    -74.013f,
+                                    -10.3369f,
+                                    71.7653f,
+                                    -108.2964f,
+                                    -1.746f,
+                                    93.1133f,
+                                    -25.8063f,
+                                    -33.0234f
+                                },
+                                {
+                                    -74.013f, 10.3369f, -25.8063f, 33.0234f, -1.746f, -93.1133f,
+                                    71.7653f, 108.2964f
+                                }
+                            });
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
+
+    @Test
+    public void testIfft2() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray array =
+                    manager.create(
+                            new float[][] {
+                                {1f, 6.6f, 4.315f, 2.0f},
+                                {16.9f, 6.697f, 2.399f, 67.9f},
+                                {0f, 5f, 67.09f, 9.87f}
+                            });
+            long[] sizes = {3, 4};
+            long[] axes = {0, 1};
+            NDArray fft2 = array.fft2(sizes, axes);
+            NDArray actual = fft2.ifft2(sizes, axes).real();
+            NDArray expected = array.toType(DataType.COMPLEX64, true).real();
+            Assertions.assertAlmostEquals(expected, actual);
+        }
+    }
 }