Allow for negative reduction axes.

PiperOrigin-RevId: 693082679
ROCm · Nov 4, 2024 · d670c15 · d670c15
1 parent eec20f5
commit d670c15
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 24 deletions.
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
@@ -4621,24 +4621,11 @@ class Subgraph {
     }
 
     if (subgraph != nullptr) {
-      std::array<size_t, XNN_MAX_TENSOR_DIMS> reduction_axes;
+      std::array<int64_t, XNN_MAX_TENSOR_DIMS> reduction_axes;
       for (int i = 0; i < num_reduction_axes; ++i) {
-        if (axes_data[i] < 0) {
-          if (IsDynamicTensor(&input_tensor)) {
-            TF_LITE_MAYBE_KERNEL_LOG(
-                logging_context,
-                "unable to delegate %s node #%d: negative reduction axis only "
-                "supported for non-dynamic input tensors",
-                EnumNameBuiltinOperator(tflite_operator), node_index);
-            return kTfLiteError;
-          }
-          reduction_axes[i] = axes_data[i] + NumDimensions(&input_tensor);
-        } else {
-          reduction_axes[i] = axes_data[i];
-        }
+        reduction_axes[i] = axes_data[i];
       }
-      std::sort(&reduction_axes[0], &reduction_axes[num_reduction_axes]);
-      if (xnn_define_static_reduce(
+      if (xnn_define_static_reduce_v2(
               subgraph, reduce_operator, num_reduction_axes,
               reduction_axes.data(),
               /*input_id=*/input_output_tensors.at(node->inputs->data[0]),

diff --git a/tensorflow/lite/tools/cmake/modules/cpuinfo.cmake b/tensorflow/lite/tools/cmake/modules/cpuinfo.cmake
@@ -23,7 +23,7 @@ OverridableFetchContent_Declare(
   cpuinfo
   GIT_REPOSITORY https://github.com/pytorch/cpuinfo
   # Sync with tensorflow/workspace2.bzl
-  GIT_TAG 1e83a2fdd3102f65c6f1fb602c1b320486218a99
+  GIT_TAG dff2616ddd49122b63abcf44d2c097483b77f861
   GIT_PROGRESS TRUE
   SOURCE_DIR "${CMAKE_BINARY_DIR}/cpuinfo"
 )

diff --git a/tensorflow/lite/tools/cmake/modules/xnnpack.cmake b/tensorflow/lite/tools/cmake/modules/xnnpack.cmake
@@ -23,7 +23,7 @@ OverridableFetchContent_Declare(
   xnnpack
   GIT_REPOSITORY https://github.com/google/XNNPACK
   # Sync with tensorflow/workspace2.bzl
-  GIT_TAG 743f95f0c34b02d6d2cdb9e87da21caffe9c668f
+  GIT_TAG d62f1ebbc4682371985fc7b948af5835634f2660
   GIT_PROGRESS TRUE
   PREFIX "${CMAKE_BINARY_DIR}"
   SOURCE_DIR "${CMAKE_BINARY_DIR}/xnnpack"

diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl
@@ -150,9 +150,9 @@ def _tf_repositories():
     # LINT.IfChange
     tf_http_archive(
         name = "XNNPACK",
-        sha256 = "bfedea7d94d4b7953a857868b63eda27a2d8206c79dd0b0456d4150cc43bf825",
-        strip_prefix = "XNNPACK-743f95f0c34b02d6d2cdb9e87da21caffe9c668f",
-        urls = tf_mirror_urls("https://github.com/google/XNNPACK/archive/743f95f0c34b02d6d2cdb9e87da21caffe9c668f.zip"),
+        sha256 = "b371299f5fbc96b965ba1089e327e0836affa66305994a7e53d846912d4335bf",
+        strip_prefix = "XNNPACK-d62f1ebbc4682371985fc7b948af5835634f2660",
+        urls = tf_mirror_urls("https://github.com/google/XNNPACK/archive/d62f1ebbc4682371985fc7b948af5835634f2660.zip"),
     )
     # LINT.ThenChange(//tensorflow/lite/tools/cmake/modules/xnnpack.cmake)
 
@@ -180,10 +180,10 @@ def _tf_repositories():
 
     tf_http_archive(
         name = "cpuinfo",
-        sha256 = "ca31f17a86e4db01b5fc05efa1807ddc84c02ba4611464b67e185e8210bf096b",
-        strip_prefix = "cpuinfo-1e83a2fdd3102f65c6f1fb602c1b320486218a99",
+        sha256 = "a57f2b93704fe49ad37fa3ef5ce8208443d6b589b89653341332bcf51aa6dde3",
+        strip_prefix = "cpuinfo-dff2616ddd49122b63abcf44d2c097483b77f861",
         patch_file = ["//third_party/cpuinfo:cpuinfo_ppc64le_support.patch"],
-        urls = tf_mirror_urls("https://github.com/pytorch/cpuinfo/archive/1e83a2fdd3102f65c6f1fb602c1b320486218a99.zip"),
+        urls = tf_mirror_urls("https://github.com/pytorch/cpuinfo/archive/dff2616ddd49122b63abcf44d2c097483b77f861.zip"),
     )
 
     tf_http_archive(