Repository: deepjavalibrary/djl
Pull request #3337 (2024-07-13 14:38:51)

NdArray ifft, rfft, irfft, conj implementation for PyTorch engine.
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/.github/workflows/continuous.yml b/.github/workflows/continuous.yml
index c9d67a542..e5bbc6850 100644
--- a/.github/workflows/continuous.yml
+++ b/.github/workflows/continuous.yml
@@ -87,7 +87,7 @@ jobs:
         run: |
           ./gradlew -Dai.djl.default_engine=OnnxRuntime :integration:clean :integration:test
       - name: Build with Gradle
-        run: ./gradlew clean build :jacoco:testCodeCoverageReport
+        run: ./gradlew build :jacoco:testCodeCoverageReport
       - name: Upload test results
         uses: actions/upload-artifact@v3
         if: always()
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
index a84b2bc03..92840065d 100644
--- a/.github/workflows/native_jni_s3_pytorch.yml
+++ b/.github/workflows/native_jni_s3_pytorch.yml
@@ -272,7 +272,12 @@ jobs:
   build-pytorch-jni-aarch64:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: [ self-hosted, aarch64 ]
-    container: amazonlinux:2
+    container:
+      image: amazonlinux:2
+      env:
+        JAVA_HOME: /usr/lib/jvm/java-17-amazon-corretto.aarch64
+        AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+        AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
     timeout-minutes: 30
     needs: create-aarch64-runner
     steps:
@@ -288,22 +293,15 @@ jobs:
           ln -sf /usr/bin/gcc10-ar /usr/bin/ar
           ln -sf /usr/bin/cmake3 /usr/bin/cmake
           pip3 install awscli --upgrade
-      - uses: actions/checkout@v3
+      - uses: taiki-e/checkout-action@v1
       - name: Release JNI prep
         run: |
-          export JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto.aarch64
           export PATH=$PATH:$JAVA_HOME/bin
           PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
           export PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)}
           echo $PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
-      - name: Configure AWS Credentials
-        uses: aws-actions/configure-aws-credentials@v2
-        with:
-          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
-          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-          aws-region: us-east-2
       - name: Copy files to S3 with the AWS CLI
         run: |
           PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
diff --git a/README.md b/README.md
index 8a7d1de17..3ed6c465f 100644
--- a/README.md
+++ b/README.md
@@ -85,14 +85,15 @@ The following pseudocode demonstrates running training:
 
 ## Release Notes
 
+* [0.29.0](https://github.com/deepjavalibrary/djl/releases/tag/v0.29.0) ([Code](https://github.com/deepjavalibrary/djl/tree/v0.29.0))
 * [0.28.0](https://github.com/deepjavalibrary/djl/releases/tag/v0.28.0) ([Code](https://github.com/deepjavalibrary/djl/tree/v0.28.0))
 * [0.27.0](https://github.com/deepjavalibrary/djl/releases/tag/v0.27.0) ([Code](https://github.com/deepjavalibrary/djl/tree/v0.27.0))
 * [0.26.0](https://github.com/deepjavalibrary/djl/releases/tag/v0.26.0) ([Code](https://github.com/deepjavalibrary/djl/tree/v0.26.0))
 * [0.25.0](https://github.com/deepjavalibrary/djl/releases/tag/v0.25.0) ([Code](https://github.com/deepjavalibrary/djl/tree/v0.25.0))
 * [0.24.0](https://github.com/deepjavalibrary/djl/releases/tag/v0.24.0) ([Code](https://github.com/deepjavalibrary/djl/tree/v0.24.0))
-* [+24 releases](https://github.com/deepjavalibrary/djl/releases)
+* [+25 releases](https://github.com/deepjavalibrary/djl/releases)
 
-The release of DJL 0.29.0 is planned for July 2024.
+The release of DJL 0.30.0 is planned for September 2024.
 
 ## Building From Source
 
diff --git a/android/pytorch-native/README.md b/android/pytorch-native/README.md
index 3d6752487..b5a1d9602 100644
--- a/android/pytorch-native/README.md
+++ b/android/pytorch-native/README.md
@@ -124,7 +124,7 @@ cd ..
 ./gradlew compileAndroidJNI -Ppt_version=${PYTORCH_VERSION}
 ```
 
-`jnilib/0.28.0/android` folder will be created after build, and shared library will be uploaded to S3 in CI build
+`jnilib/0.29.0/android` folder will be created after build, and shared library will be uploaded to S3 in CI build
 
 ## Build PyTorch android library (.aar) and publish to Sonatype snapshot repo
 
@@ -138,7 +138,7 @@ cd ../../../android
 
 # To avoid download jni from S3, manually copy them
 mkdir -p pytorch-native/jnilib
-cp -r ../engines/pytorch/pytorch-native/jnilib/0.28.0/android/* pytorch-native/jnilib
+cp -r ../engines/pytorch/pytorch-native/jnilib/0.29.0/android/* pytorch-native/jnilib
 
 ./gradlew :pytorch-native:assemble
 # publish to local maven repo (~/.m2 folder)
diff --git a/api/README.md b/api/README.md
index ff4878fb8..514465929 100644
--- a/api/README.md
+++ b/api/README.md
@@ -35,7 +35,7 @@ You can pull the DJL API from the central Maven repository by including the foll
 <dependency>
     <groupId>ai.djl</groupId>
     <artifactId>api</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
@@ -45,7 +45,7 @@ For testing the current nightly build, use the following:
 <dependency>
     <groupId>ai.djl</groupId>
     <artifactId>api</artifactId>
-    <version>0.29.0-SNAPSHOT</version>
+    <version>0.30.0-SNAPSHOT</version>
 </dependency>
 ```
 
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index abbe7df97..d9227a3c3 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -3375,6 +3375,69 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray fft(long length, long axis);
 
+    /**
+     * Computes the one dimensional inverse discrete Fourier transform.
+     *
+     * @param length Length of the transformed axis of the output.
+     * @param axis Axis over which to compute the IFFT.
+     * @return The truncated or zero-padded input, transformed along the axis indicated by axis, or
+     *     the last one if axis is not specified.
+     */
+    NDArray ifft(long length, long axis);
+
+    /**
+     * Computes the one dimensional inverse 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 ifft(long length) {
+        return ifft(length, -1);
+    }
+
+    /**
+     * Computes the one dimensional Fourier transform of real-valued input.
+     *
+     * @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 rfft(long length) {
+        return rfft(length, -1);
+    }
+
+    /**
+     * Computes the one dimensional Fourier transform of real-valued input.
+     *
+     * @param length Length of the transformed axis of the output.
+     * @param axis Axis over which to compute the FFT.
+     * @return The truncated or transformed along the axis indicated by axis, or the last one if
+     *     axis is not specified.
+     */
+    NDArray rfft(long length, long axis);
+
+    /**
+     * Computes the one dimensional inverse Fourier transform of real-valued input.
+     *
+     * @param length Length of the transformed axis of the output.
+     * @param axis Axis over which to compute the IRFFT.
+     * @return The truncated or transformed along the axis indicated by axis, or the last one if
+     *     axis is not specified.
+     */
+    NDArray irfft(long length, long axis);
+
+    /**
+     * Computes the one dimensional inverse Fourier transform of real-valued input.
+     *
+     * @param length Length of the transformed axis of the output.
+     * @return The truncated or transformed along the axis indicated by axis, or the last one if
+     *     axis is not specified.
+     */
+    default NDArray irfft(long length) {
+        return irfft(length, -1);
+    }
+
     /**
      * Computes the Short Time Fourier Transform (STFT).
      *
@@ -5404,4 +5467,12 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @return tje real NDArray
      */
     NDArray real();
+
+    /**
+     * Conjugate complex array.
+     *
+     * @return Returns a view of input with a flipped conjugate bit. If input has a non-complex
+     *     type, this function just returns input.
+     */
+    NDArray conj();
 }
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 9a4ad8db9..27c362c88 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -900,6 +900,24 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray rfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray irfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray stft(
@@ -1254,6 +1272,12 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray conj() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArrayEx getNDArrayInternal() {
diff --git a/basicdataset/README.md b/basicdataset/README.md
index 592d57e1f..f312659f1 100644
--- a/basicdataset/README.md
+++ b/basicdataset/README.md
@@ -29,7 +29,7 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl</groupId>
     <artifactId>basicdataset</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
diff --git a/bom/README.md b/bom/README.md
index 194917493..534a95132 100644
--- a/bom/README.md
+++ b/bom/README.md
@@ -22,7 +22,7 @@ will need to mention the type as pom and the scope as import) as the following:
         <dependency>
             <groupId>ai.djl</groupId>
             <artifactId>bom</artifactId>
-            <version>0.28.0</version>
+            <version>0.29.0</version>
             <type>pom</type>
             <scope>import</scope>
         </dependency>
@@ -38,7 +38,7 @@ will need to mention the type as pom and the scope as import) as the following:
         <dependency>
             <groupId>ai.djl</groupId>
             <artifactId>bom</artifactId>
-            <version>0.28.0</version>
+            <version>0.29.0</version>
             <type>pom</type>
             <scope>import</scope>
         </dependency>
@@ -65,7 +65,7 @@ will need to mention the type as pom and the scope as import) as the following:
 - First you need add BOM into your build.gradle file as the following:
 
 ```
-    implementation platform("ai.djl:bom:0.28.0")
+    implementation platform("ai.djl:bom:0.29.0")
 ```
 
 - Then you import the desired DJL modules into to you pom.xml file (no version is needed):
diff --git a/buildSrc/src/main/kotlin/ai/djl/release.gradle.kts b/buildSrc/src/main/kotlin/ai/djl/release.gradle.kts
index 96eff83af..a8f349066 100644
--- a/buildSrc/src/main/kotlin/ai/djl/release.gradle.kts
+++ b/buildSrc/src/main/kotlin/ai/djl/release.gradle.kts
@@ -51,10 +51,11 @@ tasks {
             for (file in collection) {
                 file.text = file.text.replace("/${previousVersion}/", "/${djlVersion}/")
                     .replace(">${previousVersion}<", ">${djlVersion}<")
+                    .replace(">${previousVersion}-SNAPSHOT<", ">${djlVersion}-SNAPSHOT<")
                     .replace("\"${previousVersion}\"", "\"${djlVersion}\"")
-                    .replace("-${previousVersion}", "-${djlVersion}")
+                    .replace(" ${previousVersion}`", " ${djlVersion}`")
                     .replace("_${previousVersion}", "_${djlVersion}")
-                    .replace(":${previousVersion}([\"\\\\\\n])", ":${djlVersion}\$1")
+                    .replace(Regex(":${previousVersion}([\"\\\\\\n])"), ":${djlVersion}\$1")
             }
         }
     }
diff --git a/djl-zero/README.md b/djl-zero/README.md
index 74267e9ee..dfc1ff9ef 100644
--- a/djl-zero/README.md
+++ b/djl-zero/README.md
@@ -49,6 +49,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl</groupId>
     <artifactId>djl-zero</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/docker/spark/Dockerfile b/docker/spark/Dockerfile
index c33e3b185..47c1f460a 100644
--- a/docker/spark/Dockerfile
+++ b/docker/spark/Dockerfile
@@ -13,14 +13,14 @@ FROM 314815235551.dkr.ecr.us-east-2.amazonaws.com/sagemaker-spark-processing:3.5
 LABEL maintainer="djl-dev@amazon.com"
 
 # Install dependencies
-ARG DJL_VERSION=0.28.0
+ARG DJL_VERSION=0.29.0
 ARG JNA_VERSION=5.14.0
-ARG JAVACV_VERSION=1.5.9
-ARG JAVACPP_VERSION=1.5.9
-ARG FFMPEG_VERSION=6.0-1.5.9
-ARG TENSORFLOW_CORE_VERSION=0.5.0
+ARG JAVACV_VERSION=1.5.10
+ARG JAVACPP_VERSION=1.5.10
+ARG FFMPEG_VERSION=6.1.1-1.5.10
+ARG TENSORFLOW_CORE_VERSION=1.0.0-rc.1
 ARG PROTOBUF_VERSION=3.25.3
-ARG PYTORCH_VERSION=2.2.2
+ARG PYTORCH_VERSION=2.3.1
 
 RUN pip3 install --no-cache-dir torch==${PYTORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu
 RUN pip3 install --no-cache-dir pillow pandas numpy pyarrow librosa 'transformers[torch]'
diff --git a/docs/development/example_dataset.md b/docs/development/example_dataset.md
index 8ee5f9890..80cf517b1 100644
--- a/docs/development/example_dataset.md
+++ b/docs/development/example_dataset.md
@@ -24,8 +24,8 @@ api group: 'org.apache.commons', name: 'commons-csv', version: '1.7'
 In order to extend the dataset, the following dependencies are required:
 
 ```
-api "ai.djl:api:0.28.0"
-api "ai.djl:basicdataset:0.28.0"
+api "ai.djl:api:0.29.0"
+api "ai.djl:basicdataset:0.29.0"
 ```
 
 There are four parts we need to implement for CSVDataset.
diff --git a/docs/get.md b/docs/get.md
index 2c6e8b999..d1e3f328f 100644
--- a/docs/get.md
+++ b/docs/get.md
@@ -99,7 +99,7 @@ dependencies {
     implementation platform("ai.djl:bom:<UPCOMING VERSION>-SNAPSHOT")
 }
 ```
-Currently, the `<UPCOMING VERSION> = 0.28.0`.
+Currently, the `<UPCOMING VERSION> = 0.29.0`.
 This snapshot version is the same as the custom DJL repository. 
 
 You also need to change directory to `djl/bom`. Then build and publish it to maven local same as what was done in `djl`.
diff --git a/docs/hybrid_engine.md b/docs/hybrid_engine.md
index 132205f17..d63340aa4 100644
--- a/docs/hybrid_engine.md
+++ b/docs/hybrid_engine.md
@@ -21,17 +21,17 @@ to run in a hybrid mode:
 To use it along with Apache MXNet for additional API support, add the following two dependencies:
 
 ```
-runtimeOnly "ai.djl.mxnet:mxnet-engine:0.28.0"
+runtimeOnly "ai.djl.mxnet:mxnet-engine:0.29.0"
 ```
 
 You can also use PyTorch or TensorFlow Engine as the supplemental engine by adding their corresponding dependencies.
 
 ```
-runtimeOnly "ai.djl.pytorch:pytorch-engine:0.28.0"
+runtimeOnly "ai.djl.pytorch:pytorch-engine:0.29.0"
 ```
 
 ```
-runtimeOnly "ai.djl.tensorflow:tensorflow-engine:0.28.0"
+runtimeOnly "ai.djl.tensorflow:tensorflow-engine:0.29.0"
 ```
 
 ## How Hybrid works
diff --git a/docs/load_model.md b/docs/load_model.md
index 11c632170..25af7c8c5 100644
--- a/docs/load_model.md
+++ b/docs/load_model.md
@@ -181,7 +181,7 @@ Here is a few tips you can use to help you debug model loading issue:
 See [here](development/configure_logging.md#configure-logging-level) for how to enable debug log
 
 #### List models programmatically in your code
-You can use [ModelZoo.listModels()](https://javadoc.io/static/ai.djl/api/0.28.0/ai/djl/repository/zoo/ModelZoo.html#listModels--) API to query available models.
+You can use [ModelZoo.listModels()](https://javadoc.io/static/ai.djl/api/0.29.0/ai/djl/repository/zoo/ModelZoo.html#listModels--) API to query available models.
 
 #### List available models using DJL command line
 
diff --git a/engines/ml/lightgbm/README.md b/engines/ml/lightgbm/README.md
index 6744996dd..3500f9546 100644
--- a/engines/ml/lightgbm/README.md
+++ b/engines/ml/lightgbm/README.md
@@ -36,13 +36,13 @@ LightGBM can only run on top of the Linux/Mac/Windows machine using x86_64.
 ## Installation
 You can pull the LightGBM engine from the central Maven repository by including the following dependency:
 
-- ai.djl.ml.lightgbm:lightgbm:0.28.0
+- ai.djl.ml.lightgbm:lightgbm:0.29.0
 
 ```xml
 <dependency>
     <groupId>ai.djl.ml.lightgbm</groupId>
     <artifactId>lightgbm</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
diff --git a/engines/ml/xgboost/README.md b/engines/ml/xgboost/README.md
index d72e8d535..2a1fdf22e 100644
--- a/engines/ml/xgboost/README.md
+++ b/engines/ml/xgboost/README.md
@@ -37,13 +37,13 @@ XGBoost can only run on top of the Linux/Mac machine. User can build from source
 ## Installation
 You can pull the XGBoost engine from the central Maven repository by including the following dependency:
 
-- ai.djl.ml.xgboost:xgboost:0.28.0
+- ai.djl.ml.xgboost:xgboost:0.29.0
 
 ```xml
 <dependency>
     <groupId>ai.djl.ml.xgboost</groupId>
     <artifactId>xgboost</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
diff --git a/engines/mxnet/mxnet-engine/README.md b/engines/mxnet/mxnet-engine/README.md
index fede00b2c..2a1b26417 100644
--- a/engines/mxnet/mxnet-engine/README.md
+++ b/engines/mxnet/mxnet-engine/README.md
@@ -7,7 +7,7 @@ This module contains the Deep Java Library (DJL) EngineProvider for Apache MXNet
 We don't recommend that developers use classes in this module directly. Use of these classes
 will couple your code with Apache MXNet and make switching between engines difficult. Even so,
 developers are not restricted from using engine-specific features. For more information,
-see [NDManager#invoke()](https://javadoc.io/static/ai.djl/api/0.28.0/ai/djl/ndarray/NDManager.html#invoke-java.lang.String-ai.djl.ndarray.NDArray:A-ai.djl.ndarray.NDArray:A-ai.djl.util.PairList-).
+see [NDManager#invoke()](https://javadoc.io/static/ai.djl/api/0.29.0/ai/djl/ndarray/NDManager.html#invoke-java.lang.String-ai.djl.ndarray.NDArray:A-ai.djl.ndarray.NDArray:A-ai.djl.util.PairList-).
 
 ## Documentation
 
@@ -33,7 +33,7 @@ You can pull the MXNet engine from the central Maven repository by including the
 <dependency>
     <groupId>ai.djl.mxnet</groupId>
     <artifactId>mxnet-engine</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
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 52fac203d..f6d2c0aa0 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
@@ -1152,6 +1152,24 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft(long length, long axis) {
+        return null;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray irfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray rfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray stft(
@@ -1679,6 +1697,12 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray conj() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArrayEx getNDArrayInternal() {
diff --git a/engines/mxnet/mxnet-model-zoo/README.md b/engines/mxnet/mxnet-model-zoo/README.md
index 6302cf737..8f45796ec 100644
--- a/engines/mxnet/mxnet-model-zoo/README.md
+++ b/engines/mxnet/mxnet-model-zoo/README.md
@@ -27,7 +27,7 @@ You can pull the MXNet engine from the central Maven repository by including the
 <dependency>
     <groupId>ai.djl.mxnet</groupId>
     <artifactId>mxnet-model-zoo</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
diff --git a/engines/onnxruntime/onnxruntime-android/README.md b/engines/onnxruntime/onnxruntime-android/README.md
index 565fb5e88..b3f7d0a38 100644
--- a/engines/onnxruntime/onnxruntime-android/README.md
+++ b/engines/onnxruntime/onnxruntime-android/README.md
@@ -12,7 +12,7 @@ You can pull the ONNX Runtime for Android from the central Maven repository by i
 <dependency>
     <groupId>ai.djl.android</groupId>
     <artifactId>onnxruntime</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
diff --git a/engines/onnxruntime/onnxruntime-engine/README.md b/engines/onnxruntime/onnxruntime-engine/README.md
index e11838707..e75f892a3 100644
--- a/engines/onnxruntime/onnxruntime-engine/README.md
+++ b/engines/onnxruntime/onnxruntime-engine/README.md
@@ -37,13 +37,13 @@ for the official ONNX Runtime project.
 ## Installation
 You can pull the ONNX Runtime engine from the central Maven repository by including the following dependency:
 
-- ai.djl.onnxruntime:onnxruntime-engine:0.28.0
+- ai.djl.onnxruntime:onnxruntime-engine:0.29.0
 
 ```xml
 <dependency>
     <groupId>ai.djl.onnxruntime</groupId>
     <artifactId>onnxruntime-engine</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -61,7 +61,7 @@ Maven:
 <dependency>
     <groupId>ai.djl.onnxruntime</groupId>
     <artifactId>onnxruntime-engine</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
     <exclusions>
         <exclusion>
@@ -73,7 +73,7 @@ Maven:
 <dependency>
     <groupId>com.microsoft.onnxruntime</groupId>
     <artifactId>onnxruntime_gpu</artifactId>
-    <version>1.17.1</version>
+    <version>1.18.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -81,10 +81,10 @@ Maven:
 Gradle:
 
 ```groovy
-implementation("ai.djl.onnxruntime:onnxruntime-engine:0.28.0") {
+implementation("ai.djl.onnxruntime:onnxruntime-engine:0.29.0") {
     exclude group: "com.microsoft.onnxruntime", module: "onnxruntime"
 }
-implementation "com.microsoft.onnxruntime:onnxruntime_gpu:1.17.1"
+implementation "com.microsoft.onnxruntime:onnxruntime_gpu:1.18.0"
 ```
 
 #### Enable TensorRT execution
diff --git a/engines/pytorch/pytorch-engine/README.md b/engines/pytorch/pytorch-engine/README.md
index 80779067c..a6fd6f0a9 100644
--- a/engines/pytorch/pytorch-engine/README.md
+++ b/engines/pytorch/pytorch-engine/README.md
@@ -24,13 +24,13 @@ The javadocs output is built in the `build/doc/javadoc` folder.
 ## Installation
 You can pull the PyTorch engine from the central Maven repository by including the following dependency:
 
-- ai.djl.pytorch:pytorch-engine:0.28.0
+- ai.djl.pytorch:pytorch-engine:0.29.0
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-engine</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -46,6 +46,7 @@ The following table illustrates which pytorch version that DJL supports:
 
 | PyTorch engine version | PyTorch native library version            |
 |------------------------|-------------------------------------------|
+| pytorch-engine:0.29.0  | 1.13.1, 2.1.2, 2.2.2, **2.3.1**           |
 | pytorch-engine:0.28.0  | 1.13.1, 2.1.2, **2.2.2**                  |
 | pytorch-engine:0.27.0  | 1.13.1, **2.1.1**                         |
 | pytorch-engine:0.26.0  | 1.13.1, 2.0.1, **2.1.1**                  |
@@ -115,21 +116,21 @@ export PYTORCH_FLAVOR=cpu
 ### macOS
 For macOS, you can use the following library:
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cpu:2.1.1:osx-x86_64
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cpu:2.3.1:osx-x86_64
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>osx-x86_64</classifier>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -139,21 +140,21 @@ For macOS, you can use the following library:
 ### macOS M1
 For macOS M1, you can use the following library:
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cpu:2.2.2:osx-aarch64
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cpu:2.3.1:osx-aarch64
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>osx-aarch64</classifier>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -164,29 +165,29 @@ installed on your GPU machine, you can use one of the following library:
 
 #### Linux GPU
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cu121:2.2.2:linux-x86_64 - CUDA 12.1
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cu121:2.3.1:linux-x86_64 - CUDA 12.1
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cu121</artifactId>
     <classifier>linux-x86_64</classifier>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
 
 ### Linux CPU
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cpu:2.2.2:linux-x86_64
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cpu:2.3.1:linux-x86_64
 
 ```xml
 <dependency>
@@ -194,20 +195,20 @@ installed on your GPU machine, you can use one of the following library:
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>linux-x86_64</classifier>
     <scope>runtime</scope>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
 
 ### For aarch64 build
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.2.2:linux-aarch64
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.3.1:linux-aarch64
 
 ```xml
 <dependency>
@@ -215,12 +216,12 @@ installed on your GPU machine, you can use one of the following library:
     <artifactId>pytorch-native-cpu-precxx11</artifactId>
     <classifier>linux-aarch64</classifier>
     <scope>runtime</scope>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -230,22 +231,22 @@ installed on your GPU machine, you can use one of the following library:
 We also provide packages for the system like CentOS 7/Ubuntu 14.04 with GLIBC >= 2.17.
 All the package were built with GCC 7, we provided a newer `libstdc++.so.6.24` in the package that contains `CXXABI_1.3.9` to use the package successfully.
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cu121-precxx11:2.2.2:linux-x86_64 - CUDA 12.1
-- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.2.2:linux-x86_64   - CPU
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cu121-precxx11:2.3.1:linux-x86_64 - CUDA 12.1
+- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.3.1:linux-x86_64   - CPU
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cu121-precxx11</artifactId>
     <classifier>linux-x86_64</classifier>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -255,13 +256,13 @@ All the package were built with GCC 7, we provided a newer `libstdc++.so.6.24` i
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cpu-precxx11</artifactId>
     <classifier>linux-x86_64</classifier>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -276,29 +277,29 @@ For the Windows platform, you can choose between CPU and GPU.
 
 #### Windows GPU
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cu121:2.2.2:win-x86_64 - CUDA 12.1
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cu121:2.3.1:win-x86_64 - CUDA 12.1
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cu121</artifactId>
     <classifier>win-x86_64</classifier>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
 
 ### Windows CPU
 
-- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
-- ai.djl.pytorch:pytorch-native-cpu:2.2.2:win-x86_64
+- ai.djl.pytorch:pytorch-jni:2.3.1-0.29.0
+- ai.djl.pytorch:pytorch-native-cpu:2.3.1:win-x86_64
 
 ```xml
 <dependency>
@@ -306,12 +307,12 @@ For the Windows platform, you can choose between CPU and GPU.
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>win-x86_64</classifier>
     <scope>runtime</scope>
-    <version>2.2.2</version>
+    <version>2.3.1</version>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.2.2-0.28.0</version>
+    <version>2.3.1-0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
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 f33c2b247..1edd61226 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
@@ -1101,6 +1101,24 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         return JniUtils.fft(this, length, axis);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray rfft(long length, long axis) {
+        return JniUtils.rfft(this, length, axis);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft(long length, long axis) {
+        return JniUtils.ifft(this, length, axis);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray irfft(long length, long axis) {
+        return JniUtils.irfft(this, length, axis);
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray stft(
@@ -1628,6 +1646,12 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         return JniUtils.real(this);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray conj() {
+        return JniUtils.conj(this);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDArrayEx getNDArrayInternal() {
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 53e4c76fc..b80c096cb 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
@@ -1017,6 +1017,24 @@ public final class JniUtils {
                 PyTorchLibrary.LIB.torchFft(ndArray.getHandle(), length, axis));
     }
 
+    public static PtNDArray ifft(PtNDArray ndArray, long length, long axis) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchIfft(ndArray.getHandle(), length, axis));
+    }
+
+    public static PtNDArray rfft(PtNDArray ndArray, long length, long axis) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchRfft(ndArray.getHandle(), length, axis));
+    }
+
+    public static PtNDArray irfft(PtNDArray ndArray, long length, long axis) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchIrfft(ndArray.getHandle(), length, axis));
+    }
+
     public static PtNDArray stft(
             PtNDArray ndArray,
             long nFft,
@@ -1068,6 +1086,10 @@ public final class JniUtils {
         return new PtNDArray(ndArray.getManager(), handle);
     }
 
+    public static PtNDArray conj(PtNDArray ndArray) {
+        return new PtNDArray(ndArray.getManager(), PyTorchLibrary.LIB.conj(ndArray.getHandle()));
+    }
+
     public static PtNDArray abs(PtNDArray ndArray) {
         return new PtNDArray(
                 ndArray.getManager(), PyTorchLibrary.LIB.torchAbs(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 3c03cd8cc..5c1fd290f 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
@@ -270,6 +270,12 @@ final class PyTorchLibrary {
 
     native long torchFft(long handle, long length, long axis);
 
+    native long torchIfft(long handle, long length, long axis);
+
+    native long torchRfft(long handle, long length, long axis);
+
+    native long torchIrfft(long handle, long length, long axis);
+
     native long torchStft(
             long handle,
             long nFft,
@@ -287,6 +293,8 @@ final class PyTorchLibrary {
 
     native long torchViewAsComplex(long handle);
 
+    native long conj(long handle);
+
     native long[] torchSplit(long handle, long size, long dim);
 
     native long[] torchSplit(long handle, long[] indices, long dim);
diff --git a/engines/pytorch/pytorch-model-zoo/README.md b/engines/pytorch/pytorch-model-zoo/README.md
index bdb7444de..d58658b61 100644
--- a/engines/pytorch/pytorch-model-zoo/README.md
+++ b/engines/pytorch/pytorch-model-zoo/README.md
@@ -25,7 +25,7 @@ You can pull the PyTorch engine from the central Maven repository by including t
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-model-zoo</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
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 08932098d..48ae78db5 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,33 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFft(
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIfft(
+    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_ifft(*tensor_ptr, jn, jaxis));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchRfft(
+    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_rfft(*tensor_ptr, jn, jaxis));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIrfft(
+    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_irfft(*tensor_ptr, jn, jaxis));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  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()
@@ -90,6 +117,14 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchViewAsReal(
 #endif
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_conj(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::conj(*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) {
 #ifdef V1_11_X
diff --git a/engines/tensorflow/tensorflow-api/README.md b/engines/tensorflow/tensorflow-api/README.md
index a70a8943a..204f11630 100644
--- a/engines/tensorflow/tensorflow-api/README.md
+++ b/engines/tensorflow/tensorflow-api/README.md
@@ -16,6 +16,6 @@ You can pull the TensorFlow core java API from the central Maven repository by i
 <dependency>
     <groupId>ai.djl.tensorflow</groupId>
     <artifactId>tensorflow-api</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/engines/tensorflow/tensorflow-engine/README.md b/engines/tensorflow/tensorflow-engine/README.md
index 41ba43e0f..1d7934452 100644
--- a/engines/tensorflow/tensorflow-engine/README.md
+++ b/engines/tensorflow/tensorflow-engine/README.md
@@ -28,13 +28,13 @@ The javadocs output is built in the `build/doc/javadoc` folder.
 
 You can pull the TensorFlow engine from the central Maven repository by including the following dependency:
 
-- ai.djl.tensorflow:tensorflow-engine:0.28.0
+- ai.djl.tensorflow:tensorflow-engine:0.29.0
 
 ```xml
 <dependency>
     <groupId>ai.djl.tensorflow</groupId>
     <artifactId>tensorflow-engine</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
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 363d18ca4..88df43498 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
@@ -544,6 +544,12 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray conj() {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray add(Number n) {
@@ -1173,6 +1179,24 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray rfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray irfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray stft(
diff --git a/engines/tensorflow/tensorflow-model-zoo/README.md b/engines/tensorflow/tensorflow-model-zoo/README.md
index e45ef4c07..b11da5d63 100644
--- a/engines/tensorflow/tensorflow-model-zoo/README.md
+++ b/engines/tensorflow/tensorflow-model-zoo/README.md
@@ -26,7 +26,7 @@ from the central Maven repository by including the following dependency:
 <dependency>
     <groupId>ai.djl.tensorflow</groupId>
     <artifactId>tensorflow-model-zoo</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
diff --git a/engines/tensorrt/README.md b/engines/tensorrt/README.md
index 8efefa799..7c1471eeb 100644
--- a/engines/tensorrt/README.md
+++ b/engines/tensorrt/README.md
@@ -28,13 +28,13 @@ The javadocs output is generated in the `build/doc/javadoc` folder.
 ## Installation
 You can pull the TensorRT engine from the central Maven repository by including the following dependency:
 
-- ai.djl.tensorrt:tensorrt:0.28.0
+- ai.djl.tensorrt:tensorrt:0.29.0
 
 ```xml
 <dependency>
     <groupId>ai.djl.tensorrt</groupId>
     <artifactId>tensorrt</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
diff --git a/examples/pom.xml b/examples/pom.xml
index c29ef32bc..e699f6900 100644
--- a/examples/pom.xml
+++ b/examples/pom.xml
@@ -5,12 +5,12 @@
 
     <groupId>ai.djl</groupId>
     <artifactId>examples</artifactId>
-    <version>0.29.0-SNAPSHOT</version>
+    <version>0.30.0-SNAPSHOT</version>
 
     <properties>
         <maven.compiler.source>11</maven.compiler.source>
         <maven.compiler.target>11</maven.compiler.target>
-        <djl.version>0.29.0-SNAPSHOT</djl.version>
+        <djl.version>0.30.0-SNAPSHOT</djl.version>
         <exec.mainClass>ai.djl.examples.inference.cv.ObjectDetection</exec.mainClass>
     </properties>
 
diff --git a/extensions/audio/README.md b/extensions/audio/README.md
index ffc404378..508069411 100644
--- a/extensions/audio/README.md
+++ b/extensions/audio/README.md
@@ -23,6 +23,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.audio</groupId>
     <artifactId>audio</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/extensions/aws-ai/README.md b/extensions/aws-ai/README.md
index 063d484dd..175de8871 100644
--- a/extensions/aws-ai/README.md
+++ b/extensions/aws-ai/README.md
@@ -58,6 +58,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.aws</groupId>
     <artifactId>aws-ai</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/extensions/fasttext/README.md b/extensions/fasttext/README.md
index 7276e3a98..875583d0d 100644
--- a/extensions/fasttext/README.md
+++ b/extensions/fasttext/README.md
@@ -34,7 +34,7 @@ You can pull the fastText engine from the central Maven repository by including
 <dependency>
     <groupId>ai.djl.fasttext</groupId>
     <artifactId>fasttext-engine</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
diff --git a/extensions/hadoop/README.md b/extensions/hadoop/README.md
index f9db9724b..591cadd5a 100644
--- a/extensions/hadoop/README.md
+++ b/extensions/hadoop/README.md
@@ -52,6 +52,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.hadoop</groupId>
     <artifactId>hadoop</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/extensions/opencv/README.md b/extensions/opencv/README.md
index db715f83c..5e656341a 100644
--- a/extensions/opencv/README.md
+++ b/extensions/opencv/README.md
@@ -23,6 +23,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.opencv</groupId>
     <artifactId>opencv</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/extensions/sentencepiece/README.md b/extensions/sentencepiece/README.md
index 78afbf3c4..0b599d059 100644
--- a/extensions/sentencepiece/README.md
+++ b/extensions/sentencepiece/README.md
@@ -23,6 +23,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.sentencepiece</groupId>
     <artifactId>sentencepiece</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/extensions/spark/README.md b/extensions/spark/README.md
index 0f384fd5c..3b52b9f30 100644
--- a/extensions/spark/README.md
+++ b/extensions/spark/README.md
@@ -34,7 +34,7 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.spark</groupId>
     <artifactId>spark_2.12</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
diff --git a/extensions/tablesaw/README.md b/extensions/tablesaw/README.md
index e49e51e57..eb1925637 100644
--- a/extensions/tablesaw/README.md
+++ b/extensions/tablesaw/README.md
@@ -25,6 +25,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.tablesaw</groupId>
     <artifactId>tablesaw</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/extensions/timeseries/README.md b/extensions/timeseries/README.md
index 7c7927793..0cd38b8b4 100644
--- a/extensions/timeseries/README.md
+++ b/extensions/timeseries/README.md
@@ -245,6 +245,6 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.timeseries</groupId>
     <artifactId>timeseries</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
diff --git a/extensions/timeseries/docs/forecast_with_M5_data.md b/extensions/timeseries/docs/forecast_with_M5_data.md
index c0e05b1b8..64edac0dc 100644
--- a/extensions/timeseries/docs/forecast_with_M5_data.md
+++ b/extensions/timeseries/docs/forecast_with_M5_data.md
@@ -56,7 +56,7 @@ repositories {
 }
 dependencies {
   implementation "org.apache.logging.log4j:log4j-slf4j-impl:2.17.1"
-  implementation platform("ai.djl:bom:0.28.0")
+  implementation platform("ai.djl:bom:0.29.0")
   implementation "ai.djl:api"
   implementation "ai.djl.timeseries"
   runtimeOnly "ai.djl.mxnet:mxnet-engine"
diff --git a/extensions/tokenizers/README.md b/extensions/tokenizers/README.md
index b65c94aae..d6086d0b0 100644
--- a/extensions/tokenizers/README.md
+++ b/extensions/tokenizers/README.md
@@ -23,7 +23,7 @@ You can pull the module from the central Maven repository by including the follo
 <dependency>
     <groupId>ai.djl.huggingface</groupId>
     <artifactId>tokenizers</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 
diff --git a/extensions/tokenizers/build.sh b/extensions/tokenizers/build.sh
index 58c4df4f0..f04f22aed 100755
--- a/extensions/tokenizers/build.sh
+++ b/extensions/tokenizers/build.sh
@@ -32,7 +32,7 @@ RUST_MANIFEST=rust/Cargo.toml
 if [[ "$FLAVOR" = "cpu"* ]]; then
   cargo build --manifest-path $RUST_MANIFEST --release
 elif [[ "$FLAVOR" = "cu"* && "$FLAVOR" > "cu121" ]]; then
-  cargo build --manifest-path $RUST_MANIFEST --release --features cuda,cublaslt,flash-attn
+  cargo build --manifest-path $RUST_MANIFEST --release --features cuda,flash-attn
 else
   cargo build --manifest-path $RUST_MANIFEST --release
 fi
diff --git a/extensions/tokenizers/rust/Cargo.toml b/extensions/tokenizers/rust/Cargo.toml
index 24957157a..507ae35a4 100644
--- a/extensions/tokenizers/rust/Cargo.toml
+++ b/extensions/tokenizers/rust/Cargo.toml
@@ -6,11 +6,13 @@ edition = "2021"
 
 [dependencies]
 jni = "0.21.1"
-candle = { version = "0.5.1", package = "candle-core" }
-candle-nn = "0.5.1"
-candle-transformers = "0.5.1"
-candle-flash-attn = { version = "0.5.1", optional = true }
-cudarc = { version = "0.11.6", default-features = false, features = [ "cublaslt", "f16" ], optional = true }
+candle = { version = "*", package = "candle-core" }
+candle-nn = { version = "*" }
+candle-transformers = { version = "*" }
+candle-flash-attn = { version = "*", optional = true }
+candle-cublaslt = { git = "https://github.com/huggingface/candle-cublaslt", rev = "cf789b7dd6d4abb19b03b9556442f94f0588b4a0", optional = true }
+candle-layer-norm = { git = "https://github.com/xyang16/candle-layer-norm", rev = "e574de6a7f88bafbede8edf9ee43170c6a8ce51a", optional = true }
+candle-rotary = { git = "https://github.com/huggingface/candle-rotary", rev = "0a718a0856569a92f3112e64f10d07e4447822e8", optional = true }
 tokenizers = { path = "../tokenizers/tokenizers", version = "*", features = ["http"] }
 half = "2.4.0"
 tracing = "0.1.40"
@@ -19,13 +21,19 @@ thiserror = "1.0.58"
 serde = { version = "1.0.198", features = ["serde_derive"] }
 serde_json = "1.0.116"
 
+[patch.crates-io]
+cudarc = { git = "https://github.com/coreylowman/cudarc", rev = "c388e724af93a3e8fbe484f5ded2d8b3c1badd8e" }
+candle = { git = "https://github.com/huggingface/candle", rev = "f76bb7794aa8659c5023797979a3392cdfc01f32", package = "candle-core" }
+candle-nn = { git = "https://github.com/huggingface/candle", rev = "f76bb7794aa8659c5023797979a3392cdfc01f32", package = "candle-nn" }
+candle-transformers = { git = "https://github.com/huggingface/candle", rev = "f76bb7794aa8659c5023797979a3392cdfc01f32", package = "candle-transformers" }
+candle-flash-attn = { git = "https://github.com/huggingface/candle", rev = "f76bb7794aa8659c5023797979a3392cdfc01f32", package = "candle-flash-attn" }
+
 [target.'cfg(target_os = "linux")'.dependencies]
 openssl = { version = "0.10", features = ["vendored"] }
 
 [lib]
-crate_type = ["cdylib"]
+crate-type = ["cdylib"]
 
 [features]
-cuda = ["candle/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+cuda = ["candle/cuda", "candle-nn/cuda", "candle-transformers/cuda", "dep:candle-cublaslt", "dep:candle-layer-norm", "dep:candle-rotary"]
 flash-attn = ["cuda", "candle-transformers/flash-attn", "dep:candle-flash-attn"]
-cublaslt = ["cudarc/cublaslt"]
diff --git a/extensions/tokenizers/rust/src/layers/cublaslt.rs b/extensions/tokenizers/rust/src/layers/cublaslt.rs
index 4fe5431e4..e35601e84 100644
--- a/extensions/tokenizers/rust/src/layers/cublaslt.rs
+++ b/extensions/tokenizers/rust/src/layers/cublaslt.rs
@@ -1,962 +1,13 @@
-#[cfg(feature = "cuda")]
-pub use cudarc::cublaslt::Activation;
-
 use crate::layers::HiddenAct;
-use candle::backend::BackendStorage;
-use candle::{CpuStorage, Device, Layout, Result, Shape, Storage, Tensor};
-
-#[cfg(feature = "cuda")]
-use candle::cuda_backend::WrapErr;
-use half::{bf16, f16};
-use std::sync::Arc;
+use candle::{Device, Result, Tensor};
 use std::sync::Once;
 
-use std::ffi::c_int;
-
 #[cfg(feature = "cuda")]
-use cudarc::cublaslt::{CudaBlasLT, Matmul, MatmulConfig};
+use candle_cublaslt::{fused_batch_matmul, fused_matmul, Activation, CublasLt};
 
 static INIT: Once = Once::new();
 static mut CUBLASLT: Option<CublasLtWrapper> = None;
 
-#[derive(Debug, Clone)]
-#[cfg(feature = "cuda")]
-pub struct CublasLt(Arc<CudaBlasLT>);
-
-#[cfg(feature = "cuda")]
-impl CublasLt {
-    pub fn new(device: &Device) -> Result<Self> {
-        let dev = match &*device {
-            Device::Cuda(d) => d,
-            _ => candle::bail!("`device` must be a `cuda` device"),
-        };
-
-        let inner = CudaBlasLT::new(dev.cuda_device()).unwrap();
-
-        Ok(Self(Arc::new(inner)))
-    }
-}
-
-#[cfg(feature = "cuda")]
-pub struct CublasLTMatmul {
-    pub cublaslt: Arc<CudaBlasLT>,
-    pub act: Option<Activation>,
-    pub c: Option<Tensor>,
-    pub alpha: Option<f32>,
-    pub beta: Option<f32>,
-}
-
-#[cfg(feature = "cuda")]
-impl CublasLTMatmul {
-    pub fn fwd_f16(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: Option<&candle::CudaStorage>,
-        bias_l: Option<&Layout>,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        let dev = a.device();
-
-        // Assume TN
-        let (m, k) = a_l.shape().dims2()?;
-
-        let (n, b_1) = b_l.shape().dims2()?;
-
-        if b_1 != k {
-            candle::bail!("This layer only supports TN layout");
-        }
-
-        let lda = k;
-        let ldb = k;
-        let ldc = m;
-
-        let out_shape = Shape::from((n, m));
-
-        let a = a.as_cuda_slice::<f16>()?.slice(a_l.start_offset()..);
-        let b = b.as_cuda_slice::<f16>()?.slice(b_l.start_offset()..);
-
-        let bias = if let (Some(bias), Some(bias_l)) = (bias, bias_l) {
-            if bias_l.shape().dims1()? != m {
-                candle::bail!("Bias does not have the correct shape");
-            }
-
-            Some(bias.as_cuda_slice::<f16>()?.slice(bias_l.start_offset()..))
-        } else {
-            None
-        };
-
-        let mut out = if let Some(c) = &self.c {
-            let (c, c_l) = c.storage_and_layout();
-            let c = match &*c {
-                Storage::Cuda(storage) => storage.as_cuda_slice::<f16>()?,
-                _ => candle::bail!("`c` must be a cuda tensor"),
-            };
-            match c_l.contiguous_offsets() {
-                Some((o1, o2)) => {
-                    if o1 != 0 {
-                        candle::bail!("`c` start offset must be 0");
-                    }
-                    if o2 != out_shape.elem_count() {
-                        candle::bail!("`c` end offset must be {}", out_shape.elem_count())
-                    }
-                }
-                None => candle::bail!("`c` has to be contiguous"),
-            };
-            if c_l.shape().dims2()? != (n, m) {
-                candle::bail!("`c` does not have the correct shape");
-            }
-
-            c.clone()
-        } else {
-            // Allocate out tensor
-            unsafe { dev.alloc::<f16>(out_shape.elem_count()).w()? }
-        };
-
-        let config = MatmulConfig {
-            transa: true,
-            transb: false,
-            m: m as u64,
-            n: n as u64,
-            k: k as u64,
-            alpha: self.alpha.unwrap_or(1.0),
-            lda: lda as i64,
-            ldb: ldb as i64,
-            beta: self.beta.unwrap_or(0.0),
-            ldc: ldc as i64,
-            stride_a: None,
-            stride_b: None,
-            stride_c: None,
-            stride_bias: None,
-            batch_size: None,
-        };
-
-        unsafe {
-            self.cublaslt
-                .matmul(config, &a, &b, &mut out, bias.as_ref(), self.act.as_ref())
-                .map_err(|e| candle::Error::Cuda(Box::new(e)))?;
-        }
-
-        let out = candle::CudaStorage::wrap_cuda_slice(out, dev.clone());
-
-        Ok((out, out_shape))
-    }
-
-    pub fn fwd_bf16(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: Option<&candle::CudaStorage>,
-        bias_l: Option<&Layout>,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        let dev = a.device();
-
-        // Assume TN
-        let (m, k) = a_l.shape().dims2()?;
-
-        let (n, b_1) = b_l.shape().dims2()?;
-
-        if b_1 != k {
-            candle::bail!("This layer only supports TN layout");
-        }
-
-        let lda = k;
-        let ldb = k;
-        let ldc = m;
-
-        let out_shape = Shape::from((n, m));
-
-        let a = a.as_cuda_slice::<bf16>()?.slice(a_l.start_offset()..);
-        let b = b.as_cuda_slice::<bf16>()?.slice(b_l.start_offset()..);
-
-        let bias = if let (Some(bias), Some(bias_l)) = (bias, bias_l) {
-            if bias_l.shape().dims1()? != m {
-                candle::bail!("Bias does not have the correct shape");
-            }
-
-            Some(bias.as_cuda_slice::<bf16>()?.slice(bias_l.start_offset()..))
-        } else {
-            None
-        };
-
-        let mut out = if let Some(c) = &self.c {
-            let (c, c_l) = c.storage_and_layout();
-            let c = match &*c {
-                Storage::Cuda(storage) => storage.as_cuda_slice::<bf16>()?,
-                _ => candle::bail!("`c` must be a cuda tensor"),
-            };
-            match c_l.contiguous_offsets() {
-                Some((o1, o2)) => {
-                    if o1 != 0 {
-                        candle::bail!("`c` start offset must be 0");
-                    }
-                    if o2 != out_shape.elem_count() {
-                        candle::bail!("`c` end offset must be {}", out_shape.elem_count())
-                    }
-                }
-                None => candle::bail!("`c` has to be contiguous"),
-            };
-            if c_l.shape().dims2()? != (n, m) {
-                candle::bail!("`c` does not have the correct shape");
-            }
-
-            c.clone()
-        } else {
-            // Allocate out tensor
-            unsafe { dev.alloc::<bf16>(out_shape.elem_count()).w()? }
-        };
-
-        let config = MatmulConfig {
-            transa: true,
-            transb: false,
-            m: m as u64,
-            n: n as u64,
-            k: k as u64,
-            alpha: self.alpha.unwrap_or(1.0),
-            lda: lda as i64,
-            ldb: ldb as i64,
-            beta: self.beta.unwrap_or(0.0),
-            ldc: ldc as i64,
-            stride_a: None,
-            stride_b: None,
-            stride_c: None,
-            stride_bias: None,
-            batch_size: None,
-        };
-
-        unsafe {
-            self.cublaslt
-                .matmul(config, &a, &b, &mut out, bias.as_ref(), self.act.as_ref())
-                .map_err(|e| candle::Error::Cuda(Box::new(e)))?;
-        }
-
-        let out = candle::CudaStorage::wrap_cuda_slice(out, dev.clone());
-
-        Ok((out, out_shape))
-    }
-
-    pub fn fwd_f32(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: Option<&candle::CudaStorage>,
-        bias_l: Option<&Layout>,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        let dev = a.device();
-
-        // Assume TN
-        let (m, k) = a_l.shape().dims2()?;
-
-        let (n, b_1) = b_l.shape().dims2()?;
-
-        if b_1 != k {
-            candle::bail!("This layer only supports TN layout");
-        }
-
-        let lda = k;
-        let ldb = k;
-        let ldc = m;
-
-        let out_shape = Shape::from((n, m));
-
-        let a = a.as_cuda_slice::<f32>()?.slice(a_l.start_offset()..);
-        let b = b.as_cuda_slice::<f32>()?.slice(b_l.start_offset()..);
-
-        let bias = if let (Some(bias), Some(bias_l)) = (bias, bias_l) {
-            if bias_l.shape().dims1()? != m {
-                candle::bail!("Bias does not have the correct shape");
-            }
-
-            Some(bias.as_cuda_slice::<f32>()?.slice(bias_l.start_offset()..))
-        } else {
-            None
-        };
-
-        let mut out = if let Some(c) = &self.c {
-            let (c, c_l) = c.storage_and_layout();
-            let c = match &*c {
-                Storage::Cuda(storage) => storage.as_cuda_slice::<f32>()?,
-                _ => candle::bail!("`c` must be a cuda tensor"),
-            };
-            match c_l.contiguous_offsets() {
-                Some((o1, o2)) => {
-                    if o1 != 0 {
-                        candle::bail!("`c` start offset must be 0");
-                    }
-                    if o2 != out_shape.elem_count() {
-                        candle::bail!("`c` end offset must be {}", out_shape.elem_count())
-                    }
-                }
-                None => candle::bail!("`c` has to be contiguous"),
-            };
-            if c_l.shape().dims2()? != (n, m) {
-                candle::bail!("`c` does not have the correct shape");
-            }
-
-            c.clone()
-        } else {
-            // Allocate out tensor
-            unsafe { dev.alloc::<f32>(out_shape.elem_count()).w()? }
-        };
-
-        let config = MatmulConfig {
-            transa: true,
-            transb: false,
-            m: m as u64,
-            n: n as u64,
-            k: k as u64,
-            alpha: self.alpha.unwrap_or(1.0),
-            lda: lda as i64,
-            ldb: ldb as i64,
-            beta: self.beta.unwrap_or(0.0),
-            ldc: ldc as i64,
-            stride_a: None,
-            stride_b: None,
-            stride_c: None,
-            stride_bias: None,
-            batch_size: None,
-        };
-
-        unsafe {
-            self.cublaslt
-                .matmul(config, &a, &b, &mut out, bias.as_ref(), self.act.as_ref())
-                .map_err(|e| candle::Error::Cuda(Box::new(e)))?;
-        }
-
-        let out = candle::CudaStorage::wrap_cuda_slice(out, dev.clone());
-
-        Ok((out, out_shape))
-    }
-}
-
-#[cfg(feature = "cuda")]
-impl candle::CustomOp2 for CublasLTMatmul {
-    fn name(&self) -> &'static str {
-        "cublaslt-matmul"
-    }
-
-    fn cpu_fwd(
-        &self,
-        _: &CpuStorage,
-        _: &Layout,
-        _: &CpuStorage,
-        _: &Layout,
-    ) -> Result<(CpuStorage, Shape)> {
-        candle::bail!("no cpu support for cublaslt-matmul")
-    }
-
-    fn cuda_fwd(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        match a.dtype() {
-            candle::DType::F16 => self.fwd_f16(a, a_l, b, b_l, None, None),
-            candle::DType::BF16 => self.fwd_bf16(a, a_l, b, b_l, None, None),
-            candle::DType::F32 => self.fwd_f32(a, a_l, b, b_l, None, None),
-            dt => candle::bail!("cublaslt-matmul is only supported for f16/bf16/f32 ({dt:?})"),
-        }
-    }
-}
-
-#[cfg(feature = "cuda")]
-impl candle::CustomOp3 for CublasLTMatmul {
-    fn name(&self) -> &'static str {
-        "cublaslt-matmul-add"
-    }
-
-    fn cpu_fwd(
-        &self,
-        _: &CpuStorage,
-        _: &Layout,
-        _: &CpuStorage,
-        _: &Layout,
-        _: &CpuStorage,
-        _: &Layout,
-    ) -> Result<(CpuStorage, Shape)> {
-        candle::bail!("no cpu support for cublaslt-matmul")
-    }
-
-    fn cuda_fwd(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: &candle::CudaStorage,
-        bias_l: &Layout,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        match a.dtype() {
-            candle::DType::F16 => self.fwd_f16(a, a_l, b, b_l, Some(bias), Some(bias_l)),
-            candle::DType::BF16 => self.fwd_bf16(a, a_l, b, b_l, Some(bias), Some(bias_l)),
-            candle::DType::F32 => self.fwd_f32(a, a_l, b, b_l, Some(bias), Some(bias_l)),
-            dt => candle::bail!("cublaslt-matmul is only supported for f16/bf16/f32 ({dt:?})"),
-        }
-    }
-}
-
-/// Fused matmul + add + Relu/Gelu activation using CublasLt
-///
-/// # Arguments
-///
-/// * `a` - Input tensor of size MxK
-/// * `b` - Input tensor of size NxK
-/// * `out` - Optional Output tensor of size NxK.
-///           If set and beta != 0, will be added to the end result of A*B before `act`
-/// * `alpha` - Optional scaling factor for A*B
-/// * `beta` - Optional scaling factor for C
-/// * `bias` - Optional bias tensor of size M
-/// * `act` - Optional Gelu or Relu activation. If set, will be added to the end result
-/// * `cublaslt` - CublasLt handle
-///
-/// The resulting tensor is of shape NxM
-#[cfg(feature = "cuda")]
-pub fn fused_matmul(
-    a: &Tensor,
-    b: &Tensor,
-    out: Option<&Tensor>,
-    alpha: Option<f32>,
-    beta: Option<f32>,
-    bias: Option<&Tensor>,
-    act: Option<Activation>,
-    cublaslt: CublasLt,
-) -> Result<Tensor> {
-    let op = CublasLTMatmul {
-        act,
-        cublaslt: cublaslt.0,
-        c: out.cloned(),
-        alpha,
-        beta,
-    };
-
-    if let Some(bias) = bias {
-        a.apply_op3(&b, &bias, op)
-    } else {
-        a.apply_op2(&b, op)
-    }
-}
-
-#[cfg(feature = "cuda")]
-pub struct CublasLTBatchMatmul {
-    pub cublaslt: Arc<CudaBlasLT>,
-    pub act: Option<Activation>,
-    pub c: Option<Tensor>,
-    pub alpha: Option<f32>,
-    pub beta: Option<f32>,
-}
-
-#[cfg(feature = "cuda")]
-impl CublasLTBatchMatmul {
-    pub fn fwd_f16(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: Option<&candle::CudaStorage>,
-        bias_l: Option<&Layout>,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        let dev = a.device();
-
-        // Assume TN
-        let (batch_size, m, k) = a_l.shape().dims3()?;
-        let (b_0, n, b_2) = b_l.shape().dims3()?;
-
-        if b_2 != k {
-            candle::bail!("This layer only supports TN layout");
-        }
-
-        if b_0 != batch_size {
-            candle::bail!("`b` must have the same batch size as `a`")
-        }
-
-        let lda = k;
-        let ldb = k;
-        let ldc = m;
-
-        let out_shape = Shape::from((batch_size, n, m));
-
-        let a = a.as_cuda_slice::<f16>()?.slice(a_l.start_offset()..);
-        let b = b.as_cuda_slice::<f16>()?.slice(b_l.start_offset()..);
-
-        let bias = if let (Some(bias), Some(bias_l)) = (bias, bias_l) {
-            if bias_l.shape().dims1()? != m {
-                candle::bail!("Bias does not have the correct shape");
-            }
-
-            Some(bias.as_cuda_slice::<f16>()?.slice(bias_l.start_offset()..))
-        } else {
-            None
-        };
-
-        let (mut out, stride_c) = if let Some(c) = &self.c {
-            let (c, c_l) = c.storage_and_layout();
-            let c = match &*c {
-                Storage::Cuda(storage) => storage.as_cuda_slice::<f16>()?,
-                _ => candle::bail!("`c` must be a cuda tensor"),
-            };
-            match c_l.contiguous_offsets() {
-                Some((o1, o2)) => {
-                    if o1 != 0 {
-                        candle::bail!("`c` start offset must be 0");
-                    }
-                    if o2 != out_shape.elem_count() {
-                        candle::bail!("`c` end offset must be {}", out_shape.elem_count())
-                    }
-                }
-                None => candle::bail!("`c` has to be contiguous"),
-            };
-
-            if c_l.shape().dims3()? != (batch_size, n, m) {
-                candle::bail!("`c` does not have the correct shape");
-            }
-
-            // Set beta to 0.0 if it is not set
-            (c.clone(), c_l.stride()[0])
-        } else {
-            // Allocate out tensor
-            (
-                unsafe { dev.alloc::<f16>(out_shape.elem_count()).w()? },
-                (n * m),
-            )
-        };
-
-        let config = MatmulConfig {
-            transa: true,
-            transb: false,
-            m: m as u64,
-            n: n as u64,
-            k: k as u64,
-            alpha: self.alpha.unwrap_or(1.0),
-            lda: lda as i64,
-            ldb: ldb as i64,
-            beta: self.beta.unwrap_or(0.0),
-            ldc: ldc as i64,
-            stride_a: Some(a_l.stride()[0] as i64),
-            stride_b: Some(b_l.stride()[0] as i64),
-            stride_c: Some(stride_c as i64),
-            stride_bias: None,
-            batch_size: Some(batch_size as c_int),
-        };
-
-        unsafe {
-            self.cublaslt
-                .matmul(config, &a, &b, &mut out, bias.as_ref(), self.act.as_ref())
-                .map_err(|e| candle::Error::Cuda(Box::new(e)))?;
-        }
-
-        let out = candle::CudaStorage::wrap_cuda_slice(out, dev.clone());
-
-        Ok((out, out_shape))
-    }
-
-    pub fn fwd_bf16(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: Option<&candle::CudaStorage>,
-        bias_l: Option<&Layout>,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        let dev = a.device();
-
-        // Assume TN
-        let (batch_size, m, k) = a_l.shape().dims3()?;
-        let (b_0, n, b_2) = b_l.shape().dims3()?;
-
-        if b_2 != k {
-            candle::bail!("This layer only supports TN layout");
-        }
-
-        if b_0 != batch_size {
-            candle::bail!("`b` must have the same batch size as `a`")
-        }
-
-        let lda = k;
-        let ldb = k;
-        let ldc = m;
-
-        let out_shape = Shape::from((batch_size, n, m));
-
-        let a = a.as_cuda_slice::<bf16>()?.slice(a_l.start_offset()..);
-        let b = b.as_cuda_slice::<bf16>()?.slice(b_l.start_offset()..);
-
-        let bias = if let (Some(bias), Some(bias_l)) = (bias, bias_l) {
-            if bias_l.shape().dims1()? != m {
-                candle::bail!("Bias does not have the correct shape");
-            }
-
-            Some(bias.as_cuda_slice::<bf16>()?.slice(bias_l.start_offset()..))
-        } else {
-            None
-        };
-
-        let (mut out, stride_c) = if let Some(c) = &self.c {
-            let (c, c_l) = c.storage_and_layout();
-            let c = match &*c {
-                Storage::Cuda(storage) => storage.as_cuda_slice::<bf16>()?,
-                _ => candle::bail!("`c` must be a cuda tensor"),
-            };
-            match c_l.contiguous_offsets() {
-                Some((o1, o2)) => {
-                    if o1 != 0 {
-                        candle::bail!("`c` start offset must be 0");
-                    }
-                    if o2 != out_shape.elem_count() {
-                        candle::bail!("`c` end offset must be {}", out_shape.elem_count())
-                    }
-                }
-                None => candle::bail!("`c` has to be contiguous"),
-            };
-
-            if c_l.shape().dims3()? != (batch_size, n, m) {
-                candle::bail!("`c` does not have the correct shape");
-            }
-
-            // Set beta to 0.0 if it is not set
-            (c.clone(), c_l.stride()[0])
-        } else {
-            // Allocate out tensor
-            (
-                unsafe { dev.alloc::<bf16>(out_shape.elem_count()).w()? },
-                (n * m),
-            )
-        };
-
-        let config = MatmulConfig {
-            transa: true,
-            transb: false,
-            m: m as u64,
-            n: n as u64,
-            k: k as u64,
-            alpha: self.alpha.unwrap_or(1.0),
-            lda: lda as i64,
-            ldb: ldb as i64,
-            beta: self.beta.unwrap_or(0.0),
-            ldc: ldc as i64,
-            stride_a: Some(a_l.stride()[0] as i64),
-            stride_b: Some(b_l.stride()[0] as i64),
-            stride_c: Some(stride_c as i64),
-            stride_bias: None,
-            batch_size: Some(batch_size as c_int),
-        };
-
-        unsafe {
-            self.cublaslt
-                .matmul(config, &a, &b, &mut out, bias.as_ref(), self.act.as_ref())
-                .map_err(|e| candle::Error::Cuda(Box::new(e)))?;
-        }
-
-        let out = candle::CudaStorage::wrap_cuda_slice(out, dev.clone());
-
-        Ok((out, out_shape))
-    }
-
-    pub fn fwd_f32(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: Option<&candle::CudaStorage>,
-        bias_l: Option<&Layout>,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        let dev = a.device();
-
-        // Assume TN
-        let (batch_size, m, k) = a_l.shape().dims3()?;
-        let (b_0, n, b_2) = b_l.shape().dims3()?;
-
-        if b_2 != k {
-            candle::bail!("This layer only supports TN layout");
-        }
-
-        if b_0 != batch_size {
-            candle::bail!("`b` must have the same batch size as `a`")
-        }
-
-        let lda = k;
-        let ldb = k;
-        let ldc = m;
-
-        let out_shape = Shape::from((batch_size, n, m));
-
-        let a = a.as_cuda_slice::<f32>()?.slice(a_l.start_offset()..);
-        let b = b.as_cuda_slice::<f32>()?.slice(b_l.start_offset()..);
-
-        let bias = if let (Some(bias), Some(bias_l)) = (bias, bias_l) {
-            if bias_l.shape().dims1()? != m {
-                candle::bail!("Bias does not have the correct shape");
-            }
-
-            Some(bias.as_cuda_slice::<f32>()?.slice(bias_l.start_offset()..))
-        } else {
-            None
-        };
-
-        let (mut out, stride_c) = if let Some(c) = &self.c {
-            let (c, c_l) = c.storage_and_layout();
-            let c = match &*c {
-                Storage::Cuda(storage) => storage.as_cuda_slice::<f32>()?,
-                _ => candle::bail!("`c` must be a cuda tensor"),
-            };
-            match c_l.contiguous_offsets() {
-                Some((o1, o2)) => {
-                    if o1 != 0 {
-                        candle::bail!("`c` start offset must be 0");
-                    }
-                    if o2 != out_shape.elem_count() {
-                        candle::bail!("`c` end offset must be {}", out_shape.elem_count())
-                    }
-                }
-                None => candle::bail!("`c` has to be contiguous"),
-            };
-
-            if c_l.shape().dims3()? != (batch_size, n, m) {
-                candle::bail!("`c` does not have the correct shape");
-            }
-
-            // Set beta to 0.0 if it is not set
-            (c.clone(), c_l.stride()[0])
-        } else {
-            // Allocate out tensor
-            (
-                unsafe { dev.alloc::<f32>(out_shape.elem_count()).w()? },
-                (n * m),
-            )
-        };
-
-        let config = MatmulConfig {
-            transa: true,
-            transb: false,
-            m: m as u64,
-            n: n as u64,
-            k: k as u64,
-            alpha: self.alpha.unwrap_or(1.0),
-            lda: lda as i64,
-            ldb: ldb as i64,
-            beta: self.beta.unwrap_or(0.0),
-            ldc: ldc as i64,
-            stride_a: Some(a_l.stride()[0] as i64),
-            stride_b: Some(b_l.stride()[0] as i64),
-            stride_c: Some(stride_c as i64),
-            stride_bias: None,
-            batch_size: Some(batch_size as c_int),
-        };
-
-        unsafe {
-            self.cublaslt
-                .matmul(config, &a, &b, &mut out, bias.as_ref(), self.act.as_ref())
-                .map_err(|e| candle::Error::Cuda(Box::new(e)))?;
-        }
-
-        let out = candle::CudaStorage::wrap_cuda_slice(out, dev.clone());
-
-        Ok((out, out_shape))
-    }
-}
-
-#[cfg(feature = "cuda")]
-impl candle::CustomOp2 for CublasLTBatchMatmul {
-    fn name(&self) -> &'static str {
-        "cublaslt-batch-matmul"
-    }
-
-    fn cpu_fwd(
-        &self,
-        _: &CpuStorage,
-        _: &Layout,
-        _: &CpuStorage,
-        _: &Layout,
-    ) -> Result<(CpuStorage, Shape)> {
-        candle::bail!("no cpu support for cublaslt-batch-matmul")
-    }
-
-    fn cuda_fwd(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        match a.dtype() {
-            candle::DType::F16 => self.fwd_f16(a, a_l, b, b_l, None, None),
-            candle::DType::BF16 => self.fwd_bf16(a, a_l, b, b_l, None, None),
-            candle::DType::F32 => self.fwd_f32(a, a_l, b, b_l, None, None),
-            dt => {
-                candle::bail!("cublaslt-batch-matmul is only supported for f16/bf16/f32 ({dt:?})")
-            }
-        }
-    }
-}
-
-#[cfg(feature = "cuda")]
-impl candle::CustomOp3 for CublasLTBatchMatmul {
-    fn name(&self) -> &'static str {
-        "cublaslt-batch-matmul-add"
-    }
-
-    fn cpu_fwd(
-        &self,
-        _: &CpuStorage,
-        _: &Layout,
-        _: &CpuStorage,
-        _: &Layout,
-        _: &CpuStorage,
-        _: &Layout,
-    ) -> Result<(CpuStorage, Shape)> {
-        candle::bail!("no cpu support for cublaslt-batch-matmul-add")
-    }
-
-    fn cuda_fwd(
-        &self,
-        a: &candle::CudaStorage,
-        a_l: &Layout,
-        b: &candle::CudaStorage,
-        b_l: &Layout,
-        bias: &candle::CudaStorage,
-        bias_l: &Layout,
-    ) -> Result<(candle::CudaStorage, Shape)> {
-        match a.dtype() {
-            candle::DType::F16 => self.fwd_f16(a, a_l, b, b_l, Some(bias), Some(bias_l)),
-            candle::DType::BF16 => self.fwd_bf16(a, a_l, b, b_l, Some(bias), Some(bias_l)),
-            candle::DType::F32 => self.fwd_f32(a, a_l, b, b_l, Some(bias), Some(bias_l)),
-            dt => candle::bail!(
-                "cublaslt-batch-matmul-add is only supported for f16/bf16/f32 ({dt:?})"
-            ),
-        }
-    }
-}
-
-/// Fused batch matmul + add + Relu/Gelu activation using CublasLt
-///
-/// # Arguments
-///
-/// * `a` - Input tensor of size BxMxK
-/// * `b` - Input tensor of size BxNxK
-/// * `out` - Optional Output tensor of size BxNxK.
-///           If set and beta != 0, will be added to the end result of A*B before `act`
-/// * `alpha` - Optional scaling factor for A*B
-/// * `beta` - Optional scaling factor for C
-/// * `bias` - Optional bias tensor of size M
-/// * `act` - Optional Gelu or Relu activation. If set, will be added to the end result
-/// * `cublaslt` - CublasLt handle
-///
-/// The resulting tensor is of shape NxM
-#[cfg(feature = "cuda")]
-pub fn fused_batch_matmul(
-    a: &Tensor,
-    b: &Tensor,
-    out: Option<&Tensor>,
-    alpha: Option<f32>,
-    beta: Option<f32>,
-    bias: Option<&Tensor>,
-    act: Option<Activation>,
-    cublaslt: CublasLt,
-) -> Result<Tensor> {
-    let op = CublasLTBatchMatmul {
-        act,
-        cublaslt: cublaslt.0,
-        c: out.cloned(),
-        alpha,
-        beta,
-    };
-
-    if let Some(bias) = bias {
-        a.apply_op3(&b, &bias, op)
-    } else {
-        a.apply_op2(&b, op)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use candle::{DType, Device};
-
-    fn to_vec2_round(t: Tensor, digits: i32) -> Result<Vec<Vec<f32>>> {
-        let b = 10f32.powi(digits);
-        let t = t.to_vec2::<f32>()?;
-        let t = t
-            .iter()
-            .map(|t| t.iter().map(|t| f32::round(t * b) / b).collect())
-            .collect();
-        Ok(t)
-    }
-
-    fn to_vec3_round(t: Tensor, digits: i32) -> Result<Vec<Vec<Vec<f32>>>> {
-        let b = 10f32.powi(digits);
-        let t = t.to_vec3::<f32>()?;
-        let t = t
-            .iter()
-            .map(|t| {
-                t.iter()
-                    .map(|t| t.iter().map(|t| f32::round(t * b) / b).collect())
-                    .collect()
-            })
-            .collect();
-        Ok(t)
-    }
-
-    #[test]
-    fn test_fused_matmul() -> Result<()> {
-        let device = Device::new_cuda(0)?;
-
-        let a = Tensor::randn(0., 1., (8, 4), &device)?.to_dtype(DType::F32)?;
-        let b = Tensor::randn(0., 1., (2, 4), &device)?.to_dtype(DType::F32)?;
-        let bias = Tensor::randn(0., 1., 8, &device)?.to_dtype(DType::F32)?;
-
-        let cublaslt = CublasLt::new(&device)?;
-
-        let res = fused_matmul(&a, &b, None, None, None, Some(&bias), None, cublaslt)?;
-        let expected = (b.matmul(&a.t()?)? + bias.broadcast_left(2)?)?;
-
-        assert_eq!(
-            to_vec2_round(res.to_dtype(DType::F32)?, 4)?,
-            to_vec2_round(expected.to_dtype(DType::F32)?, 4)?
-        );
-        Ok(())
-    }
-
-    #[test]
-    fn test_fused_batch_matmul() -> Result<()> {
-        let device = Device::new_cuda(0)?;
-
-        let a = Tensor::randn(0., 1., (3, 8, 4), &device)?.to_dtype(DType::F32)?;
-        let b = Tensor::randn(0., 1., (3, 2, 4), &device)?.to_dtype(DType::F32)?;
-        let c = Tensor::randn(0., 1., (3, 2, 8), &device)?.to_dtype(DType::F32)?;
-        let bias = Tensor::randn(0., 1., 8, &device)?.to_dtype(DType::F32)?;
-
-        let cublaslt = CublasLt::new(&device)?;
-
-        let res = fused_batch_matmul(
-            &a,
-            &b,
-            Some(&c),
-            None,
-            Some(1.0),
-            Some(&bias),
-            None,
-            cublaslt,
-        )?;
-        let expected = (b.matmul(&a.t()?)?.add(&c)? + bias.broadcast_left((3, 2))?)?;
-
-        assert_eq!(
-            to_vec3_round(res.to_dtype(DType::F32)?, 4)?,
-            to_vec3_round(expected.to_dtype(DType::F32)?, 4)?
-        );
-        Ok(())
-    }
-}
-
 pub fn get_cublas_lt_wrapper() -> Option<&'static CublasLtWrapper> {
     unsafe {
         INIT.call_once(|| {
diff --git a/extensions/tokenizers/rust/src/layers/layer_norm.rs b/extensions/tokenizers/rust/src/layers/layer_norm.rs
new file mode 100644
index 000000000..0c3605727
--- /dev/null
+++ b/extensions/tokenizers/rust/src/layers/layer_norm.rs
@@ -0,0 +1,83 @@
+use candle::{DType, Device, Result, Tensor, D};
+use candle_nn::VarBuilder;
+
+#[derive(Debug)]
+pub struct LayerNorm {
+    weight: Tensor,
+    bias: Tensor,
+    epsilon: f32,
+    span: tracing::Span,
+}
+
+impl LayerNorm {
+    pub fn load(vb: VarBuilder, hidden_size: usize, epsilon: f32) -> Result<Self> {
+        Ok(Self {
+            weight: vb
+                .get(hidden_size, "weight")
+                .or_else(|_| vb.get(hidden_size, "gamma"))?,
+            bias: vb
+                .get(hidden_size, "bias")
+                .or_else(|_| vb.get(hidden_size, "beta"))?,
+            epsilon,
+            span: tracing::span!(tracing::Level::TRACE, "layer-norm"),
+        })
+    }
+
+    pub fn forward(&self, hidden_states: &Tensor, residual: Option<&Tensor>) -> Result<Tensor> {
+        let _enter = self.span.enter();
+
+        match hidden_states.device() {
+            Device::Cpu | Device::Metal(_) => {
+                let mut hidden_states = hidden_states.clone();
+                if let Some(residual) = residual {
+                    hidden_states = hidden_states.add(residual)?;
+                }
+                let hidden_states_dtype = hidden_states.dtype();
+                let internal_dtype = match hidden_states_dtype {
+                    DType::F16 | DType::BF16 => DType::F32,
+                    d => d,
+                };
+                let hidden_size = hidden_states.dim(D::Minus1)?;
+                let hidden_states = hidden_states.to_dtype(internal_dtype)?;
+                let mean_hidden_states =
+                    (hidden_states.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
+                let hidden_states = hidden_states.broadcast_sub(&mean_hidden_states)?;
+                let norm_hidden_states =
+                    (hidden_states.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
+                let hidden_states_normed = hidden_states
+                    .broadcast_div(&(norm_hidden_states + self.epsilon as f64)?.sqrt()?)?;
+                let hidden_states = hidden_states_normed
+                    .to_dtype(hidden_states_dtype)?
+                    .broadcast_mul(&self.weight)?;
+                hidden_states.broadcast_add(&self.bias)
+            }
+            Device::Cuda(_) => {
+                #[cfg(feature = "cuda")]
+                {
+                    use candle_layer_norm::{fused_add_layer_norm, layer_norm};
+
+                    let original_shape = hidden_states.shape();
+                    let hidden_states = hidden_states.flatten_to(D::Minus2)?;
+
+                    let result = if let Some(residual) = residual {
+                        let residual = residual.flatten_to(D::Minus2)?;
+
+                        let (result, _) = fused_add_layer_norm(
+                            &hidden_states,
+                            &residual,
+                            &self.weight,
+                            Some(&self.bias),
+                            self.epsilon,
+                        )?;
+                        Ok(result)
+                    } else {
+                        layer_norm(&hidden_states, &self.weight, Some(&self.bias), self.epsilon)
+                    }?;
+                    result.reshape(original_shape)
+                }
+                #[cfg(not(feature = "cuda"))]
+                candle::bail!("`cuda` feature is not enabled")
+            }
+        }
+    }
+}
diff --git a/extensions/tokenizers/rust/src/layers/linear.rs b/extensions/tokenizers/rust/src/layers/linear.rs
index 2b67d777b..0804abe0b 100644
--- a/extensions/tokenizers/rust/src/layers/linear.rs
+++ b/extensions/tokenizers/rust/src/layers/linear.rs
@@ -8,6 +8,7 @@ use serde::Deserialize;
 pub enum HiddenAct {
     Gelu,
     Relu,
+    #[serde(alias = "silu")]
     Swiglu,
 }
 
diff --git a/extensions/tokenizers/rust/src/layers/mod.rs b/extensions/tokenizers/rust/src/layers/mod.rs
index 3f0cd444c..81f633108 100644
--- a/extensions/tokenizers/rust/src/layers/mod.rs
+++ b/extensions/tokenizers/rust/src/layers/mod.rs
@@ -1,5 +1,12 @@
 #[allow(dead_code, unused)]
 mod cublaslt;
+mod layer_norm;
 mod linear;
+#[allow(dead_code, unused)]
+mod rms_norm;
 
+pub use cublaslt::get_cublas_lt_wrapper;
+pub use layer_norm::LayerNorm;
 pub use linear::{HiddenAct, Linear};
+#[allow(unused_imports)]
+pub use rms_norm::RMSNorm;
diff --git a/extensions/tokenizers/rust/src/layers/rms_norm.rs b/extensions/tokenizers/rust/src/layers/rms_norm.rs
new file mode 100644
index 000000000..e7dab642e
--- /dev/null
+++ b/extensions/tokenizers/rust/src/layers/rms_norm.rs
@@ -0,0 +1,96 @@
+use candle::{DType, Device, Result, Tensor, D};
+use candle_nn::VarBuilder;
+
+#[derive(Debug)]
+pub struct RMSNorm {
+    weight: Tensor,
+    epsilon: f32,
+    span: tracing::Span,
+}
+
+impl RMSNorm {
+    pub fn load(vb: VarBuilder, hidden_size: usize, epsilon: f32) -> Result<Self> {
+        Ok(Self {
+            weight: vb
+                .get(hidden_size, "weight")
+                .or_else(|_| vb.get(hidden_size, "gamma"))?,
+            epsilon,
+            span: tracing::span!(tracing::Level::TRACE, "rms-norm"),
+        })
+    }
+
+    pub fn forward(
+        &self,
+        hidden_states: &Tensor,
+        residual: Option<&Tensor>,
+    ) -> Result<(Tensor, Tensor)> {
+        let _enter = self.span.enter();
+
+        match hidden_states.device() {
+            Device::Cpu | Device::Metal(_) => {
+                let mut hidden_states = hidden_states.clone();
+                let residual_add = if let Some(residual) = residual {
+                    let residual_add = hidden_states.add(residual)?;
+                    hidden_states = residual_add.clone();
+                    residual_add
+                } else {
+                    hidden_states.clone()
+                };
+
+                let hidden_states_dtype = hidden_states.dtype();
+                let internal_dtype = match hidden_states_dtype {
+                    DType::F16 | DType::BF16 => DType::F32,
+                    d => d,
+                };
+                let hidden_size = hidden_states.dim(D::Minus1)?;
+                let hidden_states = hidden_states.to_dtype(internal_dtype)?;
+                let norm_hidden_states =
+                    (hidden_states.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
+                let hidden_states_normed = hidden_states
+                    .broadcast_div(&(norm_hidden_states + self.epsilon as f64)?.sqrt()?)?;
+                Ok((
+                    hidden_states_normed
+                        .to_dtype(hidden_states_dtype)?
+                        .broadcast_mul(&self.weight)?,
+                    residual_add,
+                ))
+            }
+            Device::Cuda(_) => {
+                #[cfg(feature = "cuda")]
+                {
+                    use candle_layer_norm::{fused_add_rms_norm, rms_norm};
+
+                    let original_shape = hidden_states.shape();
+                    let hidden_states = hidden_states.flatten_to(D::Minus2)?;
+
+                    if let Some(residual) = residual {
+                        let residual = residual.flatten_to(D::Minus2)?;
+
+                        let (result, residual_add) = fused_add_rms_norm(
+                            &hidden_states,
+                            &residual,
+                            &self.weight,
+                            None,
+                            self.epsilon,
+                        )?;
+                        Ok((
+                            result.reshape(original_shape)?,
+                            residual_add.reshape(original_shape)?,
+                        ))
+                    } else {
+                        let residual_add = hidden_states.clone();
+
+                        let result = rms_norm(&hidden_states, &self.weight, None, self.epsilon)?;
+
+                        Ok((
+                            result.reshape(original_shape)?,
+                            residual_add.reshape(original_shape)?,
+                        ))
+                    }
+                }
+                #[cfg(not(feature = "cuda"))]
+                candle::bail!("`cuda` feature is not enabled")
+            }
+        }
+    }
+}
diff --git a/extensions/tokenizers/rust/src/models/bert.rs b/extensions/tokenizers/rust/src/models/bert.rs
index c241232d3..cd278180e 100644
--- a/extensions/tokenizers/rust/src/models/bert.rs
+++ b/extensions/tokenizers/rust/src/models/bert.rs
@@ -1,8 +1,7 @@
-use crate::layers::Linear;
+use crate::layers::{LayerNorm, Linear};
 use crate::models::Model;
 use candle::{Device, Result, Tensor};
 use candle_nn::{embedding, Embedding, Module, VarBuilder};
-use candle_transformers::models::with_tracing::{layer_norm, LayerNorm};
 use serde::Deserialize;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
@@ -40,6 +39,8 @@ impl HiddenActLayer {
 pub enum PositionEmbeddingType {
     #[default]
     Absolute,
+    Alibi,
+    Rope,
 }
 
 // https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/configuration_bert.py#L1
@@ -160,10 +161,10 @@ impl BertEmbeddings {
             config.hidden_size,
             vb.pp("token_type_embeddings"),
         )?;
-        let layer_norm = layer_norm(
-            config.hidden_size,
-            config.layer_norm_eps,
+        let layer_norm = LayerNorm::load(
             vb.pp("LayerNorm"),
+            config.hidden_size,
+            config.layer_norm_eps as f32,
         )?;
         Ok(Self {
             word_embeddings,
@@ -187,7 +188,7 @@ impl BertEmbeddings {
             let position_ids = Tensor::new(&position_ids[..], input_ids.device())?;
             embeddings = embeddings.broadcast_add(&position_embeddings.forward(&position_ids)?)?
         }
-        let embeddings = self.layer_norm.forward(&embeddings)?;
+        let embeddings = self.layer_norm.forward(&embeddings, None)?;
         let embeddings = self.dropout.forward(&embeddings)?;
         Ok(embeddings)
     }
@@ -300,10 +301,10 @@ struct BertSelfOutput {
 impl BertSelfOutput {
     fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
         let dense = Linear::load(vb.pp("dense"), config.hidden_size, config.hidden_size, None)?;
-        let layer_norm = layer_norm(
-            config.hidden_size,
-            config.layer_norm_eps,
+        let layer_norm = LayerNorm::load(
             vb.pp("LayerNorm"),
+            config.hidden_size,
+            config.layer_norm_eps as f32,
         )?;
         let dropout = Dropout::new(config.hidden_dropout_prob);
         Ok(Self {
@@ -318,7 +319,7 @@ impl BertSelfOutput {
         let _enter = self.span.enter();
         let hidden_states = self.dense.forward(hidden_states)?;
         let hidden_states = self.dropout.forward(&hidden_states)?;
-        self.layer_norm.forward(&(hidden_states + input_tensor)?)
+        self.layer_norm.forward(&hidden_states, Some(&input_tensor))
     }
 }
 
@@ -398,10 +399,10 @@ impl BertOutput {
             config.hidden_size,
             None,
         )?;
-        let layer_norm = layer_norm(
-            config.hidden_size,
-            config.layer_norm_eps,
+        let layer_norm = LayerNorm::load(
             vb.pp("LayerNorm"),
+            config.hidden_size,
+            config.layer_norm_eps as f32,
         )?;
         let dropout = Dropout::new(config.hidden_dropout_prob);
         Ok(Self {
@@ -416,7 +417,7 @@ impl BertOutput {
         let _enter = self.span.enter();
         let hidden_states = self.dense.forward(hidden_states)?;
         let hidden_states = self.dropout.forward(&hidden_states)?;
-        self.layer_norm.forward(&(hidden_states + input_tensor)?)
+        self.layer_norm.forward(&hidden_states, Some(&input_tensor))
     }
 }
 
diff --git a/extensions/tokenizers/rust/src/models/mod.rs b/extensions/tokenizers/rust/src/models/mod.rs
index a0b52daa4..87553f3a4 100644
--- a/extensions/tokenizers/rust/src/models/mod.rs
+++ b/extensions/tokenizers/rust/src/models/mod.rs
@@ -13,6 +13,14 @@ use jni::JNIEnv;
 use serde::Deserialize;
 use std::path::PathBuf;
 
+#[derive(Debug, PartialEq, Clone)]
+pub enum Pool {
+    Cls,
+    Mean,
+    Splade,
+    LastToken,
+}
+
 pub(crate) trait Model {
     #[allow(dead_code)]
     fn is_padded(&self) -> bool;
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
index 8799c389a..ba526b098 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
@@ -1089,6 +1089,24 @@ public class RsNDArray extends NativeResource<Long> implements NDArray {
         throw new UnsupportedOperationException("Not implemented");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray rfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray irfft(long length, long axis) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray stft(
@@ -1589,6 +1607,12 @@ public class RsNDArray extends NativeResource<Long> implements NDArray {
         return toArray(RustLibrary.real(getHandle()));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray conj() {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
     /** {@inheritDoc} */
     @Override
     public RsNDArrayEx getNDArrayInternal() {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
index 295f014c5..8d5b3f962 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
@@ -17,11 +17,14 @@ import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.CrossEncoderServingTranslator;
 import ai.djl.modality.nlp.translator.TextClassificationServingTranslator;
+import ai.djl.translate.ArgumentsUtil;
 import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
+import ai.djl.util.StringPair;
 
 import java.io.IOException;
 import java.io.Serializable;
@@ -40,6 +43,7 @@ public class TextClassificationTranslatorFactory implements TranslatorFactory, S
 
     static {
         SUPPORTED_TYPES.add(new Pair<>(String.class, Classifications.class));
+        SUPPORTED_TYPES.add(new Pair<>(StringPair.class, float[].class));
         SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
     }
 
@@ -62,6 +66,17 @@ public class TextClassificationTranslatorFactory implements TranslatorFactory, S
                             .optTokenizerPath(modelPath)
                             .optManager(model.getNDManager())
                             .build();
+            if (ArgumentsUtil.booleanValue(arguments, "reranking")) {
+                CrossEncoderTranslator translator =
+                        CrossEncoderTranslator.builder(tokenizer, arguments).build();
+                if (input == StringPair.class && output == float[].class) {
+                    return (Translator<I, O>) translator;
+                } else if (input == Input.class && output == Output.class) {
+                    return (Translator<I, O>) new CrossEncoderServingTranslator(translator);
+                }
+                throw new IllegalArgumentException("Unsupported input/output types.");
+            }
+
             TextClassificationTranslator translator =
                     TextClassificationTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == Classifications.class) {
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index cd5844d6e..5ed5cf56e 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -2,7 +2,7 @@
 format.version = "1.1"
 
 [versions]
-djl = "0.29.0"
+djl = "0.30.0"
 pytorch = "2.3.1"
 tensorflow = "2.16.1"
 tensorflowCore = "1.0.0-rc.1"
diff --git a/model-zoo/README.md b/model-zoo/README.md
index dc9a00bf3..139403002 100644
--- a/model-zoo/README.md
+++ b/model-zoo/README.md
@@ -33,7 +33,7 @@ You can pull the model zoo from the central Maven repository by including the fo
 <dependency>
     <groupId>ai.djl</groupId>
     <artifactId>model-zoo</artifactId>
-    <version>0.28.0</version>
+    <version>0.29.0</version>
 </dependency>
 ```
 


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
index 33cbd9bd5..14336182d 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
@@ -14,6 +14,7 @@ package ai.djl.huggingface.tokenizers;
 
 import ai.djl.Model;
 import ai.djl.ModelException;
+import ai.djl.huggingface.translator.TextClassificationTranslatorFactory;
 import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Input;
@@ -77,6 +78,26 @@ public class CrossEncoderTranslatorTest {
             Assert.assertEquals(res[0], 0.32456556f, 0.0001);
         }
 
+        criteria =
+                Criteria.builder()
+                        .setTypes(StringPair.class, float[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optArgument("tokenizerPath", modelDir)
+                        .optArgument("reranking", true)
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextClassificationTranslatorFactory())
+                        .build();
+
+        try (ZooModel<StringPair, float[]> model = criteria.loadModel();
+                Predictor<StringPair, float[]> predictor = model.newPredictor()) {
+            StringPair input = new StringPair(text1, text2);
+            float[] res = predictor.predict(input);
+            Assert.assertEquals(res[0], 0.32456556f, 0.0001);
+        }
+
         Criteria<Input, Output> criteria2 =
                 Criteria.builder()
                         .setTypes(Input.class, Output.class)
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 66bb136ab..214c5091b 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
@@ -20,6 +20,7 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDArrays;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.index.NDIndex;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.testing.Assertions;
@@ -1050,6 +1051,58 @@ public class NDArrayOtherOpTest {
         }
     }
 
+    @Test
+    public void testIfft() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray ele =
+                    manager.create(
+                                    new float[][] {
+                                        {15.0f, 0.0f},
+                                        {-2.5f, 3.440955f},
+                                        {-2.4999998f, 0.8122992f},
+                                        {-2.4999998f, -0.8122992f},
+                                        {-2.5f, -3.440955f}
+                                    })
+                            .complex();
+
+            NDArray result = ele.ifft(5);
+            result = result.real().get(new NDIndex(":, 0"));
+            NDArray expected = manager.create(new float[] {1, 2, 3, 4, 5});
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
+
+    @Test
+    public void testRfft() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray ele = manager.create(new float[] {1, 2, 3, 4, 5});
+            NDArray result = ele.rfft(5);
+            result = result.real().flatten();
+            NDArray expected =
+                    manager.create(
+                            new float[] {15.0f, 0.0f, -2.5f, 3.4409549f, -2.4999998f, 0.8122992f});
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
+
+    @Test
+    public void testIrfft() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray ele =
+                    manager.create(
+                                    new float[][] {
+                                        {15.0f, 0.0f},
+                                        {-2.5f, 3.4409549f},
+                                        {-2.4999998f, 0.8122992f}
+                                    })
+                            .complex();
+
+            NDArray result = ele.irfft(5);
+            NDArray expected = manager.create(new float[] {1, 2, 3, 4, 5});
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
+
     @Test
     public void testStft() {
         try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
@@ -1141,4 +1194,20 @@ public class NDArrayOtherOpTest {
             Assertions.assertAlmostEquals(expected, actual);
         }
     }
+
+    @Test
+    public void testConj() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray array = manager.create(new float[][] {{1, 2}, {-3, -4}}).complex();
+            NDArray result = array.conj().real();
+            NDArray expected = manager.create(new float[][] {{1, -2}, {-3, 4}});
+            Assertions.assertAlmostEquals(result, expected);
+
+            // non-complex array
+            array = manager.create(new float[][] {{1, 2}, {-3, -4}});
+            result = array.conj();
+            expected = manager.create(new float[][] {{1, 2}, {-3, -4}});
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
 }