From 3a5f1ed460d6c3ca458a85021ab5f6a16827a4fb Mon Sep 17 00:00:00 2001
From: Sunita Nadampalli <nadampal@amazon.com>
Date: Fri, 14 Jun 2024 14:16:41 +0000
Subject: [PATCH 1/2] linux-arm64: add build artifact item for linux-arm64
 platform

TensorFlow v2.16.1 wheel artifacts and the jar build support are added
---
 pom.xml                                              | 12 ++++++++++++
 tensorflow-core/tensorflow-core-native/pom.xml       | 10 ++++++++++
 .../tensorflow-core-native/scripts/dist_download.sh  |  3 +++
 tensorflow-core/tensorflow-core-platform/pom.xml     |  8 +++++++-
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/pom.xml b/pom.xml
index 41d80cf24fa..5dde6698c4d 100644
--- a/pom.xml
+++ b/pom.xml
@@ -329,6 +329,18 @@
         </property>
       </activation>
     </profile>
+    <profile>
+      <id>linux-arm64</id>
+      <activation>
+        <os>
+          <name>linux</name>
+          <arch>aarch64</arch>
+        </os>
+        <property>
+          <name>!javacpp.platform.extension</name>
+        </property>
+      </activation>
+    </profile>
     <profile>
       <id>macosx</id>
       <activation>
diff --git a/tensorflow-core/tensorflow-core-native/pom.xml b/tensorflow-core/tensorflow-core-native/pom.xml
index fff8f9c5756..d385b75e46d 100644
--- a/tensorflow-core/tensorflow-core-native/pom.xml
+++ b/tensorflow-core/tensorflow-core-native/pom.xml
@@ -131,6 +131,12 @@
                       <version>${project.version}</version>
                       <classifier>${javacpp.platform.windows-x86_64}</classifier>
                     </artifactItem>
+                    <artifactItem>
+                      <groupId>${project.groupId}</groupId>
+                      <artifactId>${project.artifactId}</artifactId>
+                      <version>${project.version}</version>
+                      <classifier>${javacpp.platform.linux-arm64}</classifier>
+                    </artifactItem>
                   </artifactItems>
                 </configuration>
               </execution>
@@ -169,6 +175,10 @@
                       <file>${project.build.directory}/${project.artifactId}-${project.version}-${javacpp.platform.windows-x86_64}.jar</file>
                       <classifier>${javacpp.platform.windows-x86_64}</classifier>
                     </artifact>
+                    <artifact>
+                      <file>${project.build.directory}/${project.artifactId}-${project.version}-${javacpp.platform.linux-arm64}.jar</file>
+                      <classifier>${javacpp.platform.linux-arm64}</classifier>
+                    </artifact>
                   </artifacts>
                 </configuration>
               </execution>
diff --git a/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh b/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh
index c5b10d50d48..39235d137be 100755
--- a/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh
+++ b/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh
@@ -20,6 +20,9 @@ case ${PLATFORM:-} in
     WHEEL_URL='https://files.pythonhosted.org/packages/e0/36/6278e4e7e69a90c00e0f82944d8f2713dd85a69d1add455d9e50446837ab/tensorflow_intel-2.16.1-cp311-cp311-win_amd64.whl'
     CLIB_URL='https://storage.googleapis.com/tensorflow/versions/2.16.1/libtensorflow-cpu-windows-x86_64.zip'
     ;;
+  'linux-arm64')
+    WHEEL_URL='https://files.pythonhosted.org/packages/41/ab/e5386c722548df2043c1eaadc431ea3ba0ee42a66b3af7f8013bbbacecd3/tensorflow-2.16.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl'
+    ;;
   *)
     echo "TensorFlow distribution for ${PLATFORM} is not supported for download"
     exit 1;
diff --git a/tensorflow-core/tensorflow-core-platform/pom.xml b/tensorflow-core/tensorflow-core-platform/pom.xml
index 5ed680a8bf2..0d43c7fc4b3 100644
--- a/tensorflow-core/tensorflow-core-platform/pom.xml
+++ b/tensorflow-core/tensorflow-core-platform/pom.xml
@@ -43,6 +43,12 @@
       <version>${project.version}</version>
       <classifier>linux-x86_64</classifier>
     </dependency>
+    <dependency>
+      <groupId>org.tensorflow</groupId>
+      <artifactId>tensorflow-core-native</artifactId>
+      <version>${project.version}</version>
+      <classifier>linux-arm64</classifier>
+    </dependency>
     <dependency>
       <groupId>org.tensorflow</groupId>
       <artifactId>tensorflow-core-native</artifactId>
@@ -73,7 +79,7 @@
             <configuration>
               <archive>
                 <manifestEntries>
-                  <Class-Path>tensorflow-core-api.jar tensorflow-core-native.jar tensorflow-core-native-linux-x86_64.jar tensorflow-core-native-macosx-arm64.jar tensorflow-core-native-macosx-x86_64.jar tensorflow-core-native-windows-x86_64.jar</Class-Path>
+                  <Class-Path>tensorflow-core-api.jar tensorflow-core-native.jar tensorflow-core-native-linux-x86_64.jar tensorflow-core-native-macosx-arm64.jar tensorflow-core-native-macosx-x86_64.jar tensorflow-core-native-windows-x86_64.jar tensorflow-core-native-linux-arm64.jar</Class-Path>
                   <Automatic-Module-Name>${java.module.name}</Automatic-Module-Name>
                 </manifestEntries>
               </archive>

From c081bad93b8431475f27b642f22b8712c515c32b Mon Sep 17 00:00:00 2001
From: Sunita Nadampalli <nadampal@amazon.com>
Date: Wed, 19 Jun 2024 13:50:08 +0000
Subject: [PATCH 2/2] linux-arm64 fixes to package the right dependencies into
 tensorflow native jar

---
 .../tensorflow-core-native/scripts/dist_download.sh         | 6 ++++++
 .../org/tensorflow/internal/c_api/presets/tensorflow.java   | 3 +++
 2 files changed, 9 insertions(+)

diff --git a/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh b/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh
index 39235d137be..20068238fc3 100755
--- a/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh
+++ b/tensorflow-core/tensorflow-core-native/scripts/dist_download.sh
@@ -37,6 +37,9 @@ if [[ -n "$WHEEL_URL" ]]; then
     curl -L $WHEEL_URL --output 'tensorflow.whl'
   fi
   yes | unzip -q -u 'tensorflow.whl' # use 'yes' because for some reasons -u does not work on Windows
+  if [[ "$PLATFORM" == "linux-arm64" ]]; then
+    cp $DOWNLOAD_FOLDER/tensorflow.libs/*  $DOWNLOAD_FOLDER/tensorflow/
+  fi
 fi
 
 if [[ -n "$CLIB_URL" ]]; then
@@ -51,6 +54,9 @@ cd tensorflow
 if [[ "$PLATFORM" =~ "linux" ]]; then
   ln -fs libtensorflow_cc.so.2 libtensorflow_cc.so
   ln -fs libtensorflow_framework.so.2 libtensorflow_framework.so
+  if [[ "$PLATFORM" == "linux-arm64" ]]; then
+    ln -fs libomp-*.so.5 libomp.so
+  fi
 elif [[ "$PLATFORM" =~ "macosx" ]]; then
   ln -fs libtensorflow_cc.2.dylib libtensorflow_cc.dylib
   ln -fs libtensorflow_framework.2.dylib libtensorflow_framework.dylib
diff --git a/tensorflow-core/tensorflow-core-native/src/main/java/org/tensorflow/internal/c_api/presets/tensorflow.java b/tensorflow-core/tensorflow-core-native/src/main/java/org/tensorflow/internal/c_api/presets/tensorflow.java
index ff9f02709e3..eefd6893597 100644
--- a/tensorflow-core/tensorflow-core-native/src/main/java/org/tensorflow/internal/c_api/presets/tensorflow.java
+++ b/tensorflow-core/tensorflow-core-native/src/main/java/org/tensorflow/internal/c_api/presets/tensorflow.java
@@ -62,6 +62,9 @@
           },
           link = {"tensorflow_cc@.2", "tensorflow_framework@.2"},
           resource = {"LICENSE", "THIRD_PARTY_TF_JNI_LICENSES"}),
+      @Platform(
+          value = {"linux-arm64"},
+          link = {"tensorflow_cc@.2", "tensorflow_framework@.2", "omp@.5"}),
       @Platform(
           value = "windows",
           preload = {