chore: import upstream snapshot with attribution

tensorflow/lite/delegates/ynnpack/BUILD

@@ -0,0 +1,48 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "ynnpack_delegate",
+    srcs = [
+        "copy.cc",
+        "copy.h",
+        "dot.cc",
+        "dot.h",
+        "elementwise.cc",
+        "elementwise.h",
+        "pooling.cc",
+        "pooling.h",
+        "reduction.cc",
+        "reduction.h",
+        "softmax.cc",
+        "softmax.h",
+        "utils.cc",
+        "utils.h",
+        "ynnpack_delegate.cc",
+    ],
+    hdrs = ["ynnpack_delegate.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    linkstatic = True,
+    deps = [
+        "//tensorflow/lite:builtin_ops",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/delegates/utils:simple_delegate",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/types:half",
+        "@XNNPACK//ynnpack",
+        "@XNNPACK//ynnpack/composites",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/types:span",
+        "@slinky//slinky/base:thread_pool",
+        "@slinky//slinky/base:thread_pool_impl",
+    ],
+)

tensorflow/lite/delegates/ynnpack/README.md

@@ -0,0 +1,49 @@
+# YNNPACK Delegate for LiteRT
+
+> [!WARNING]
+> The YNNPACK delegate is **experimental** and under active development. Expect
+> bugs and performance issues when using it.
+
+The YNNPACK delegate allows LiteRT (formerly TensorFlow Lite) to offload
+supported operators to YNNPACK.
+
+YNNPACK aims to provide great flexibility with good performance.
+
+## Delegate Provider Options
+
+When using LiteRT tooling (e.g., benchmarks, evaluation tools) that link the
+`ynnpack_delegate_provider`, the following command-line flags are exposed to
+configure the YNNPACK delegate:
+
+### Core Options
+
+*   **`--use_ynnpack=true|false`** (default: `false`):
+    Explicitly apply the YNNPACK delegate to the model.
+
+*   **`--num_threads=N`** (default: `0` or `1` depending on tool):
+    The number of threads to use for execution. Note that YNNPACK will only use
+    a thread pool for `num_threads > 1`. A value of `0` or `1` disables the
+    thread pool (single-threaded execution).
+
+### YNNPACK Specific Options
+
+*   **`--ynnpack_static_shape=true|false`** (default: `false`):
+    Make input shapes static instead of dynamic. Enabling this may improve
+    execution (`Invoke`) performance by allowing YNNPACK to optimize for fixed
+    shapes, but it makes model reshaping (`ResizeInputTensor`) much more
+    expensive.
+
+*   **`--ynnpack_fast_math=true|false`** (default: `false`):
+    Enable `YNN_FLAG_FAST_MATH`. This allows YNNPACK to use faster but
+    potentially less precise mathematical approximations.
+
+*   **`--ynnpack_consistent_arithmetic=true|false`** (default: `false`):
+    Enable `YNN_FLAG_CONSISTENT_ARITHMETIC`. YNNPACK will attempt to provide
+    numerically consistent results for all hardware the **same build** of
+    YNNPACK runs on. It does not guarantee consistency across different builds
+    (which means it does not guarantee consistency across different platforms,
+    which are necessarily different builds).
+
+*   **`--ynnpack_no_excess_precision=true|false`** (default: `false`):
+    Enable `YNN_FLAG_NO_EXCESS_PRECISION`. YNNPACK will not promote tensors to
+    higher precision as a performance optimization.

tensorflow/lite/delegates/ynnpack/copy.h

@@ -0,0 +1,117 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_COPY_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_COPY_H_
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsTransposeSupported(const TfLiteRegistration* registration,
+                                  const TfLiteNode* node,
+                                  TfLiteContext* context);
+
+TfLiteStatus IsSliceSupported(const TfLiteRegistration* registration,
+                              const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus IsExpandDimsSupported(const TfLiteRegistration* registration,
+                                   const TfLiteNode* node,
+                                   TfLiteContext* context);
+
+TfLiteStatus IsConcatenationSupported(const TfLiteRegistration* registration,
+                                      const TfLiteNode* node,
+                                      TfLiteContext* context);
+
+TfLiteStatus IsReshapeSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus IsPadSupported(const TfLiteRegistration* registration,
+                            const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus IsSplitSupported(const TfLiteRegistration* registration,
+                              const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus IsSpaceToDepthSupported(const TfLiteRegistration* registration,
+                                     const TfLiteNode* node,
+                                     TfLiteContext* context);
+
+TfLiteStatus IsDepthToSpaceSupported(const TfLiteRegistration* registration,
+                                     const TfLiteNode* node,
+                                     TfLiteContext* context);
+
+TfLiteStatus DefineTransposeNode(TfLiteContext* context,
+                                 ynn_subgraph_t subgraph,
+                                 TensorToValueIdMap& tensor_to_value_id,
+                                 const NodeInfo& node);
+
+TfLiteStatus DefineSliceNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                             TensorToValueIdMap& tensor_to_value_id,
+                             const NodeInfo& node);
+
+TfLiteStatus DefineConcatenationNode(TfLiteContext* context,
+                                     ynn_subgraph_t subgraph,
+                                     TensorToValueIdMap& tensor_to_value_id,
+                                     const NodeInfo& node);
+
+TfLiteStatus DefineReshapeNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                               TensorToValueIdMap& tensor_to_value_id,
+                               const NodeInfo& node);
+
+TfLiteStatus DefineExpandDimsNode(TfLiteContext* context,
+                                  ynn_subgraph_t subgraph,
+                                  TensorToValueIdMap& tensor_to_value_id,
+                                  const NodeInfo& node);
+
+TfLiteStatus DefinePadNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                           TensorToValueIdMap& tensor_to_value_id,
+                           const NodeInfo& node);
+
+TfLiteStatus DefineSplitNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                             TensorToValueIdMap& tensor_to_value_id,
+                             const NodeInfo& node);
+
+TfLiteStatus DefineSpaceToDepthNode(TfLiteContext* context,
+                                    ynn_subgraph_t subgraph,
+                                    TensorToValueIdMap& tensor_to_value_id,
+                                    const NodeInfo& node);
+
+TfLiteStatus DefineDepthToSpaceNode(TfLiteContext* context,
+                                    ynn_subgraph_t subgraph,
+                                    TensorToValueIdMap& tensor_to_value_id,
+                                    const NodeInfo& node);
+
+TfLiteStatus IsGatherSupported(const TfLiteRegistration* registration,
+                               const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus DefineGatherNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                              TensorToValueIdMap& tensor_to_value_id,
+                              const NodeInfo& node);
+
+TfLiteStatus IsGatherNdSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context);
+
+TfLiteStatus DefineGatherNdNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                                TensorToValueIdMap& tensor_to_value_id,
+                                const NodeInfo& node);
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_COPY_H_

tensorflow/lite/delegates/ynnpack/dot.cc

@@ -0,0 +1,956 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/ynnpack/dot.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <numeric>
+#include <utility>
+#include <vector>
+
+#include "ynnpack/composites/composites.h"  // from @XNNPACK
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "absl/types/span.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ynnpack {
+
+namespace {
+
+TfLiteStatus DefineQuantizedDot(
+    TfLiteContext* context, ynn_subgraph_t subgraph, int rank_a, int rank_b,
+    absl::Span<const int32_t> a_reduce_axes,
+    absl::Span<const int32_t> b_reduce_axes, uint32_t a_id, uint32_t a_scale_id,
+    uint32_t a_zp_id, uint32_t b_id, uint32_t b_scale_id, uint32_t b_zp_id,
+    uint32_t bias_id, int bias_rank, uint32_t out_scale_id, uint32_t out_zp_id,
+    bool is_per_channel, bool is_conv, ynn_type output_ynn_type,
+    uint32_t* output_id) {
+  TF_LITE_ENSURE_EQ(context, a_reduce_axes.size(), b_reduce_axes.size());
+  int num_k_dims = a_reduce_axes.size();
+  // We assume a_id and b_id are quantized (int8 or uint8).
+  // Accumulator will be int32.
+
+  int rank = rank_a - num_k_dims + 1;
+  bool is_dynamically_quantized = (out_scale_id == YNN_INVALID_VALUE_ID);
+
+  // If grouped (rank_b == 5) and per-channel, expand b_scale and b_zp to rank 3
+  // [G, 1, CO_pg].
+  if (is_conv && rank_b == 5 && is_per_channel) {
+    if (b_scale_id != YNN_INVALID_VALUE_ID) {
+      uint32_t expanded_scale_id = YNN_INVALID_VALUE_ID;
+      int32_t axes[] = {1};
+      TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_expand_dims(
+          subgraph, 1, axes, b_scale_id, &expanded_scale_id, 0));
+      b_scale_id = expanded_scale_id;
+    }
+    if (b_zp_id != YNN_INVALID_VALUE_ID) {
+      uint32_t expanded_zp_id = YNN_INVALID_VALUE_ID;
+      int32_t axes[] = {1};
+      TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_expand_dims(
+          subgraph, 1, axes, b_zp_id, &expanded_zp_id, 0));
+      b_zp_id = expanded_zp_id;
+    }
+  }
+
+  // Compute zero point and scale of the dot product.
+  uint32_t dot_zp_id = YNN_INVALID_VALUE_ID;
+  uint32_t dot_scale_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn::define_dot_quantization(
+      subgraph, num_k_dims, a_id, a_zp_id, a_scale_id, b_id, b_zp_id,
+      b_scale_id, dot_zp_id, dot_scale_id));
+
+  uint32_t accum_init_id = YNN_INVALID_VALUE_ID;
+
+  if (dot_zp_id != YNN_INVALID_VALUE_ID) {
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_unary(subgraph, ynn_unary_negate,
+                                               dot_zp_id, &accum_init_id, 0));
+  }
+
+  if (bias_id != YNN_INVALID_VALUE_ID) {
+    if (accum_init_id == YNN_INVALID_VALUE_ID) {
+      accum_init_id = bias_id;
+    } else {
+      uint32_t sub_id = YNN_INVALID_VALUE_ID;
+      uint32_t broadcasted_bias_id = bias_id;
+      TF_LITE_ENSURE_STATUS(
+          ImplementMutualBroadcasting(context, subgraph, rank, bias_rank, 0, 0,
+                                      accum_init_id, broadcasted_bias_id));
+      TF_LITE_ENSURE_YNN_STATUS(
+          ynn_define_binary(subgraph, ynn_binary_add, accum_init_id,
+                            broadcasted_bias_id, &sub_id, 0));
+      accum_init_id = sub_id;
+    }
+  }
+
+  // Now define the dot product.
+  uint32_t accum_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot(subgraph, num_k_dims, a_id, b_id,
+                                           accum_init_id, &accum_id, 0));
+
+  uint32_t accum_scale_id = dot_scale_id;
+
+  if (is_dynamically_quantized) {
+    // Dequantize accumulator directly to output.
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_dequantize(subgraph, accum_id, YNN_INVALID_VALUE_ID,
+                              accum_scale_id, ynn_type_fp32, output_id, 0));
+  } else {
+    // Dequantize accumulator to float.
+    uint32_t float_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_dequantize(subgraph, accum_id, YNN_INVALID_VALUE_ID,
+                              accum_scale_id, ynn_type_fp32, &float_id, 0));
+
+    // Quantize back to output.
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_quantize(subgraph, float_id,
+                                                  output_ynn_type, out_zp_id,
+                                                  out_scale_id, output_id, 0));
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineMatMul(TfLiteContext* context, ynn_subgraph_t subgraph,
+                          int rank_a, int rank_b, uint32_t a_id, uint32_t b_id,
+                          uint32_t bias_id, bool adj_x, bool adj_y,
+                          bool mutual_broadcast,
+                          const TfLiteTensor& input_a_tensor,
+                          const TfLiteTensor& input_b_tensor,
+                          const TfLiteTensor& output_tensor,
+                          uint32_t* output_id) {
+  bool is_input_a_quantized = IsQuantized(input_a_tensor);
+  bool is_input_b_quantized = IsQuantized(input_b_tensor);
+  bool is_output_quantized = IsQuantized(output_tensor);
+
+  bool is_quantized = is_input_a_quantized && is_output_quantized;
+  bool is_dynamically_quantized =
+      !is_input_a_quantized && is_input_b_quantized && !is_output_quantized;
+
+  bool is_per_channel = false;
+  if (is_quantized || is_dynamically_quantized) {
+    const auto* quant_params = static_cast<const TfLiteAffineQuantization*>(
+        input_b_tensor.quantization.params);
+    is_per_channel =
+        quant_params && quant_params->scale && quant_params->scale->size > 1;
+  }
+
+  uint32_t a_scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t a_zp_id = YNN_INVALID_VALUE_ID;
+  uint32_t b_scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t b_zp_id = YNN_INVALID_VALUE_ID;
+  uint32_t out_scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t out_zp_id = YNN_INVALID_VALUE_ID;
+
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+      context, subgraph, input_b_tensor, &b_scale_id, &b_zp_id));
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+      context, subgraph, output_tensor, &out_scale_id, &out_zp_id));
+
+  uint32_t current_a_id = a_id;
+  uint32_t current_b_id = b_id;
+
+  if (is_dynamically_quantized) {
+    // 1. Reduce min_max. Last axis is K (axis -1 or -2 depending on adj_x).
+    int32_t reduce_axis = adj_x ? -2 : -1;
+    uint32_t min_max_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_reduce(
+        subgraph, ynn_reduce_min_max, 1, &reduce_axis, current_a_id,
+        YNN_INVALID_VALUE_ID, &min_max_id, YNN_NODE_FLAG_KEEP_DIMS));
+
+    // 2. Define dynamic quantization params.
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_dynamic_quantization(
+        subgraph, min_max_id, ynn_type_int8, &a_zp_id, &a_scale_id, 0));
+
+    // 3. Quantize input.
+    uint32_t quantized_a_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_quantize(subgraph, current_a_id, ynn_type_int8, a_zp_id,
+                            a_scale_id, &quantized_a_id, 0));
+
+    current_a_id = quantized_a_id;
+  } else {
+    TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+        context, subgraph, input_a_tensor, &a_scale_id, &a_zp_id));
+  }
+
+  auto transpose = [&](int rank, uint32_t& val_id) -> TfLiteStatus {
+    uint32_t transposed_id = YNN_INVALID_VALUE_ID;
+    int32_t perm[YNN_MAX_TENSOR_RANK];
+    std::iota(perm, perm + rank, 0);
+    std::swap(perm[rank - 1], perm[rank - 2]);
+
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose(
+        subgraph, rank, perm, val_id, &transposed_id, 0));
+    val_id = transposed_id;
+    return kTfLiteOk;
+  };
+
+  if (adj_x) {
+    TF_LITE_ENSURE_STATUS(transpose(rank_a, current_a_id));
+    if (is_dynamically_quantized) {
+      TF_LITE_ENSURE_STATUS(transpose(rank_a, a_zp_id));
+      TF_LITE_ENSURE_STATUS(transpose(rank_a, a_scale_id));
+    }
+  }
+  if (adj_y) {
+    TF_LITE_ENSURE_STATUS(transpose(rank_b, current_b_id));
+  }
+
+  if (mutual_broadcast) {
+    TF_LITE_ENSURE_STATUS(ImplementMutualBroadcasting(
+        context, subgraph, rank_a, rank_b, /*exclude_a=*/2, /*exclude_b=*/2,
+        current_a_id, current_b_id));
+    rank_a = std::max(rank_a, rank_b);
+    rank_b = rank_a;
+  }
+
+  // Broadcast bias if present (only for float and fully quantized cases).
+  uint32_t broadcasted_bias_id = bias_id;
+
+  if (is_quantized || is_dynamically_quantized) {
+    uint32_t dot_output_id = *output_id;
+    if (is_dynamically_quantized && bias_id != YNN_INVALID_VALUE_ID) {
+      // We need intermediate output for dot before adding bias.
+      dot_output_id = YNN_INVALID_VALUE_ID;
+      TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(subgraph, ynn_type_fp32,
+                                                  rank_a, nullptr, nullptr, 0,
+                                                  &dot_output_id));
+    }
+
+    TF_LITE_ENSURE_STATUS(DefineQuantizedDot(
+        context, subgraph, rank_a, rank_b, {rank_a - 1}, {rank_b - 2},
+        current_a_id, a_scale_id, a_zp_id, current_b_id, b_scale_id, b_zp_id,
+        is_dynamically_quantized ? YNN_INVALID_VALUE_ID : broadcasted_bias_id,
+        rank_a - 1, out_scale_id, out_zp_id, is_per_channel, /*is_conv=*/false,
+        GetYnnType(output_tensor.type), &dot_output_id));
+
+    if (is_dynamically_quantized && bias_id != YNN_INVALID_VALUE_ID) {
+      TF_LITE_ENSURE_YNN_STATUS(ynn_define_binary(
+          subgraph, ynn_binary_add, dot_output_id, bias_id, output_id, 0));
+    } else {
+      *output_id = dot_output_id;
+    }
+  } else {
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_dot(subgraph, /*num_k_dims=*/1, current_a_id, current_b_id,
+                       broadcasted_bias_id, output_id, 0));
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace
+
+TfLiteStatus IsBatchMatMulSupported(const TfLiteRegistration* registration,
+                                    const TfLiteNode* node,
+                                    TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 2);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input_a = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& input_b = context->tensors[node->inputs->data[1]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, tflite::NumElements(&input_a) > 0);
+  TF_LITE_ENSURE(context, tflite::NumElements(&input_b) > 0);
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input_a));
+  TF_LITE_ENSURE(context,
+                 IsTensorSupported(input_b, /*allow_per_channel=*/true));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+
+  auto is_float_type = [](TfLiteType type) {
+    return type == kTfLiteFloat32 || type == kTfLiteFloat16 ||
+           type == kTfLiteBFloat16;
+  };
+
+  if (input_a.type == kTfLiteInt8) {
+    TF_LITE_ENSURE(context, input_b.type == kTfLiteInt8 ||
+                                input_b.type == kTfLiteInt4 ||
+                                input_b.type == kTfLiteUInt4 ||
+                                input_b.type == kTfLiteInt2);
+    TF_LITE_ENSURE(context,
+                   output.type == kTfLiteInt8 || output.type == kTfLiteInt32);
+  } else if (is_float_type(input_a.type)) {
+    if (!is_float_type(input_b.type)) {
+      TF_LITE_ENSURE(context, input_b.type == kTfLiteInt8 ||
+                                  input_b.type == kTfLiteInt4 ||
+                                  input_b.type == kTfLiteUInt4 ||
+                                  input_b.type == kTfLiteInt2);
+      const auto* params =
+          static_cast<const TfLiteBatchMatMulParams*>(node->builtin_data);
+      TF_LITE_ENSURE(context,
+                     params != nullptr && params->asymmetric_quantize_inputs);
+    }
+    TF_LITE_ENSURE(context, is_float_type(output.type));
+  } else {
+    return kTfLiteError;
+  }
+
+  if (input_b.type == kTfLiteInt4 || input_b.type == kTfLiteUInt4) {
+    TF_LITE_ENSURE(context, input_b.dims->size >= 2);
+    TF_LITE_ENSURE_EQ(context, input_b.dims->data[input_b.dims->size - 1] % 2,
+                      0);
+  } else if (input_b.type == kTfLiteInt2) {
+    TF_LITE_ENSURE(context, input_b.dims->size >= 2);
+    TF_LITE_ENSURE_EQ(context, input_b.dims->data[input_b.dims->size - 1] % 4,
+                      0);
+  }
+
+  TF_LITE_ENSURE(context, input_a.dims->size >= 2);
+  TF_LITE_ENSURE(context, input_b.dims->size >= 2);
+  TF_LITE_ENSURE(context, input_a.dims->size <= YNN_MAX_TENSOR_RANK);
+  TF_LITE_ENSURE(context, input_b.dims->size <= YNN_MAX_TENSOR_RANK);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus IsFullyConnectedSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context) {
+  TF_LITE_ENSURE(context, node->inputs->size == 2 || node->inputs->size == 3);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& weights = context->tensors[node->inputs->data[1]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, tflite::NumElements(&input) > 0);
+  TF_LITE_ENSURE(context, tflite::NumElements(&weights) > 0);
+
+  bool has_bias = node->inputs->size == 3 && node->inputs->data[2] >= 0;
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context,
+                 IsTensorSupported(weights, /*allow_per_channel=*/true));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+  if (has_bias) {
+    const TfLiteTensor& bias = context->tensors[node->inputs->data[2]];
+    TF_LITE_ENSURE(context,
+                   IsTensorSupported(bias, /*allow_per_channel=*/true));
+  }
+
+  auto is_float_type = [](TfLiteType type) {
+    return type == kTfLiteFloat32 || type == kTfLiteFloat16 ||
+           type == kTfLiteBFloat16;
+  };
+
+  if (input.type == kTfLiteInt8) {
+    TF_LITE_ENSURE(context, weights.type == kTfLiteInt8 ||
+                                weights.type == kTfLiteInt4 ||
+                                weights.type == kTfLiteUInt4 ||
+                                weights.type == kTfLiteInt2);
+    TF_LITE_ENSURE_EQ(context, output.type, kTfLiteInt8);
+    if (has_bias) {
+      const TfLiteTensor& bias = context->tensors[node->inputs->data[2]];
+      TF_LITE_ENSURE_EQ(context, bias.type, kTfLiteInt32);
+    }
+  } else if (is_float_type(input.type)) {
+    TF_LITE_ENSURE(context, is_float_type(weights.type) ||
+                                weights.type == kTfLiteInt8 ||
+                                weights.type == kTfLiteInt4 ||
+                                weights.type == kTfLiteUInt4 ||
+                                weights.type == kTfLiteInt2);
+    TF_LITE_ENSURE(context, is_float_type(output.type));
+    if (has_bias) {
+      const TfLiteTensor& bias = context->tensors[node->inputs->data[2]];
+      TF_LITE_ENSURE(context, is_float_type(bias.type));
+    }
+  } else {
+    return kTfLiteError;
+  }
+
+  if (weights.type == kTfLiteInt4 || weights.type == kTfLiteUInt4) {
+    TF_LITE_ENSURE_EQ(context, weights.dims->size, 2);
+    TF_LITE_ENSURE_EQ(context, weights.dims->data[0] % 2, 0);
+    TF_LITE_ENSURE_EQ(context, weights.dims->data[1] % 2, 0);
+  } else if (weights.type == kTfLiteInt2) {
+    TF_LITE_ENSURE_EQ(context, weights.dims->size, 2);
+    TF_LITE_ENSURE_EQ(context, weights.dims->data[0] % 4, 0);
+    TF_LITE_ENSURE_EQ(context, weights.dims->data[1] % 4, 0);
+  }
+
+  TF_LITE_ENSURE(context, input.dims->size >= 2);
+  TF_LITE_ENSURE_EQ(context, weights.dims->size, 2);
+  TF_LITE_ENSURE(context, input.dims->size <= YNN_MAX_TENSOR_RANK);
+
+  const auto* params =
+      static_cast<const TfLiteFullyConnectedParams*>(node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+  TF_LITE_ENSURE(context,
+                 IsActivationSupported(params->activation, output.type));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineBatchMatMulNode(TfLiteContext* context,
+                                   ynn_subgraph_t subgraph,
+                                   TensorToValueIdMap& tensor_to_value_id,
+                                   const NodeInfo& node) {
+  TF_LITE_ENSURE_EQ(context, node.inputs.size(), 2);
+  TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1);
+
+  int input_a_tensor_index = node.inputs[0];
+  int input_b_tensor_index = node.inputs[1];
+  int output_tensor_index = node.outputs[0];
+
+  const TfLiteTensor& input_a_tensor = context->tensors[input_a_tensor_index];
+  const TfLiteTensor& input_b_tensor = context->tensors[input_b_tensor_index];
+
+  uint32_t input_a_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input_a_tensor_index);
+  uint32_t input_b_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input_b_tensor_index);
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+
+  TF_LITE_ENSURE(context, input_a_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, input_b_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, output_val_id != YNN_INVALID_VALUE_ID);
+
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLiteBatchMatMulParams*>(tflite_node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  int rank_a = input_a_tensor.dims->size;
+  int rank_b = input_b_tensor.dims->size;
+
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+
+  TF_LITE_ENSURE_STATUS(DefineMatMul(
+      context, subgraph, rank_a, rank_b, input_a_val_id, input_b_val_id,
+      YNN_INVALID_VALUE_ID, params->adj_x, params->adj_y,
+      /*mutual_broadcast=*/true, input_a_tensor, input_b_tensor, output_tensor,
+      &output_val_id));
+
+  tensor_to_value_id[output_tensor_index] = output_val_id;
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineFullyConnectedNode(TfLiteContext* context,
+                                      ynn_subgraph_t subgraph,
+                                      TensorToValueIdMap& tensor_to_value_id,
+                                      const NodeInfo& node) {
+  const int input_tensor_index = node.inputs[0];
+  const int weights_tensor_index = node.inputs[1];
+  const int bias_tensor_index = (node.inputs.size() == 3) ? node.inputs[2] : -1;
+  const int output_tensor_index = node.outputs[0];
+
+  const TfLiteTensor& input_tensor = context->tensors[input_tensor_index];
+  const TfLiteTensor& weights_tensor = context->tensors[weights_tensor_index];
+
+  uint32_t input_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input_tensor_index);
+  uint32_t weights_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, weights_tensor_index);
+  uint32_t bias_val_id = YNN_INVALID_VALUE_ID;
+  if (bias_tensor_index != -1) {
+    bias_val_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                     bias_tensor_index);
+  }
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+
+  TF_LITE_ENSURE(context, input_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, weights_val_id != YNN_INVALID_VALUE_ID);
+  if (bias_tensor_index != -1) {
+    TF_LITE_ENSURE(context, bias_val_id != YNN_INVALID_VALUE_ID);
+  }
+  TF_LITE_ENSURE(context, output_val_id != YNN_INVALID_VALUE_ID);
+
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLiteFullyConnectedParams*>(tflite_node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  int rank_a = input_tensor.dims->size;
+  int rank_b = weights_tensor.dims->size;
+
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+
+  uint32_t reshaped_input_val_id = input_val_id;
+  if (!params->keep_num_dims) {
+    size_t input_channels = weights_tensor.dims->data[1];
+    size_t new_dims[2] = {0, input_channels};
+    reshaped_input_val_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_reshape(
+        subgraph, 2, new_dims, input_val_id, &reshaped_input_val_id, 0));
+    rank_a = 2;
+  }
+
+  // If activation is present, we must create a temporary tensor for the
+  // MatMul output, and then apply activation to it, writing to output_val_id.
+  uint32_t matmul_output_id = params->activation != kTfLiteActNone
+                                  ? YNN_INVALID_VALUE_ID
+                                  : output_val_id;
+
+  TF_LITE_ENSURE_STATUS(
+      DefineMatMul(context, subgraph, rank_a, rank_b, reshaped_input_val_id,
+                   weights_val_id, bias_val_id, /*adj_x=*/false, /*adj_y=*/true,
+                   /*mutual_broadcast=*/false, input_tensor, weights_tensor,
+                   output_tensor, &matmul_output_id));
+
+  if (params->activation != kTfLiteActNone) {
+    TF_LITE_ENSURE_STATUS(ApplyActivation(
+        context, subgraph, params->activation, matmul_output_id, output_val_id,
+        output_tensor_index, GetYnnType(output_tensor.type)));
+  }
+
+  tensor_to_value_id[output_tensor_index] = output_val_id;
+  return kTfLiteOk;
+}
+
+TfLiteStatus IsConvSupported(const TfLiteRegistration* registration,
+                             const TfLiteNode* node, TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 3);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& filter = context->tensors[node->inputs->data[1]];
+  const TfLiteTensor& bias = context->tensors[node->inputs->data[2]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context,
+                 IsTensorSupported(filter, /*allow_per_channel=*/true));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+  TF_LITE_ENSURE(context, IsTensorSupported(bias, /*allow_per_channel=*/true));
+
+  auto is_float_type = [](TfLiteType type) {
+    return type == kTfLiteFloat32 || type == kTfLiteFloat16 ||
+           type == kTfLiteBFloat16;
+  };
+
+  bool is_quantized = input.type == kTfLiteInt8 || input.type == kTfLiteUInt8;
+
+  if (is_quantized) {
+    TF_LITE_ENSURE_EQ(context, input.type, output.type);
+    TF_LITE_ENSURE_EQ(context, filter.type, kTfLiteInt8);
+    TF_LITE_ENSURE_EQ(context, bias.type, kTfLiteInt32);
+  } else if (is_float_type(input.type)) {
+    TF_LITE_ENSURE_EQ(context, input.type, output.type);
+    TF_LITE_ENSURE_EQ(context, filter.type, input.type);
+    TF_LITE_ENSURE_EQ(context, bias.type, input.type);
+  } else {
+    return kTfLiteError;
+  }
+
+  TF_LITE_ENSURE_EQ(context, input.dims->size, 4);
+  TF_LITE_ENSURE_EQ(context, filter.dims->size, 4);
+  TF_LITE_ENSURE_EQ(context, bias.dims->size, 1);
+  if (output.dims->size > 0) {
+    TF_LITE_ENSURE_EQ(context, output.dims->size, 4);
+  }
+
+  const auto* params = static_cast<const TfLiteConvParams*>(node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+  TF_LITE_ENSURE(context, params->stride_height > 0);
+  TF_LITE_ENSURE(context, params->stride_width > 0);
+  TF_LITE_ENSURE(context, params->dilation_height_factor > 0);
+  TF_LITE_ENSURE(context, params->dilation_width_factor > 0);
+
+  TF_LITE_ENSURE(context,
+                 IsActivationSupported(params->activation, output.type));
+
+  int output_channels = filter.dims->data[0];
+  int input_channels_per_group = filter.dims->data[3];
+
+  if (output.dims->size > 0) {
+    TF_LITE_ENSURE_EQ(context, output.dims->data[3], output_channels);
+  }
+  TF_LITE_ENSURE_EQ(context, bias.dims->data[0], output_channels);
+
+  int input_channels = input.dims->data[3];
+  TF_LITE_ENSURE_EQ(context, input_channels % input_channels_per_group, 0);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus IsDepthwiseConvSupported(const TfLiteRegistration* registration,
+                                      const TfLiteNode* node,
+                                      TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 3);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& filter = context->tensors[node->inputs->data[1]];
+  const TfLiteTensor& bias = context->tensors[node->inputs->data[2]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context,
+                 IsTensorSupported(filter, /*allow_per_channel=*/true));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+  TF_LITE_ENSURE(context, IsTensorSupported(bias));
+  auto is_float_type = [](TfLiteType type) {
+    return type == kTfLiteFloat32 || type == kTfLiteFloat16 ||
+           type == kTfLiteBFloat16;
+  };
+
+  bool is_quantized = input.type == kTfLiteInt8 || input.type == kTfLiteUInt8;
+
+  if (is_quantized) {
+    TF_LITE_ENSURE_EQ(context, input.type, output.type);
+    TF_LITE_ENSURE_EQ(context, filter.type, kTfLiteInt8);
+    TF_LITE_ENSURE_EQ(context, bias.type, kTfLiteInt32);
+  } else if (is_float_type(input.type)) {
+    TF_LITE_ENSURE_EQ(context, input.type, output.type);
+    TF_LITE_ENSURE_EQ(context, filter.type, input.type);
+    TF_LITE_ENSURE_EQ(context, bias.type, input.type);
+  } else {
+    return kTfLiteError;
+  }
+
+  TF_LITE_ENSURE_EQ(context, input.dims->size, 4);
+  TF_LITE_ENSURE_EQ(context, filter.dims->size, 4);
+  TF_LITE_ENSURE_EQ(context, bias.dims->size, 1);
+  if (output.dims->size > 0) {
+    TF_LITE_ENSURE_EQ(context, output.dims->size, 4);
+  }
+
+  const auto* params =
+      static_cast<const TfLiteDepthwiseConvParams*>(node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+  TF_LITE_ENSURE(context, params->stride_height > 0);
+  TF_LITE_ENSURE(context, params->stride_width > 0);
+  TF_LITE_ENSURE(context, params->dilation_height_factor > 0);
+  TF_LITE_ENSURE(context, params->dilation_width_factor > 0);
+
+  TF_LITE_ENSURE(context,
+                 IsActivationSupported(params->activation, output.type));
+
+  TF_LITE_ENSURE_EQ(context, filter.dims->data[0], 1);
+  int filter_channels = filter.dims->data[3];
+
+  int input_channels = input.dims->data[3];
+  TF_LITE_ENSURE(context, input_channels > 0);
+  TF_LITE_ENSURE_EQ(context, filter_channels % input_channels, 0);
+  int depth_multiplier = filter_channels / input_channels;
+
+  if (output.dims->size > 0) {
+    TF_LITE_ENSURE_EQ(context, output.dims->data[3],
+                      input_channels * depth_multiplier);
+  }
+  TF_LITE_ENSURE_EQ(context, bias.dims->data[0],
+                    input_channels * depth_multiplier);
+
+  if (params->depth_multiplier > 0) {
+    TF_LITE_ENSURE_EQ(context, params->depth_multiplier, depth_multiplier);
+  }
+
+  return kTfLiteOk;
+}
+
+namespace {
+
+TfLiteStatus DefineConv(TfLiteContext* context, ynn_subgraph_t subgraph,
+                        uint32_t input_id, uint32_t filter_id, uint32_t bias_id,
+                        uint32_t output_id, const TfLiteTensor& input_tensor,
+                        const TfLiteTensor& filter_tensor,
+                        const TfLiteTensor& output_tensor, int stride_height,
+                        int stride_width, int dilation_height,
+                        int dilation_width, TfLitePadding padding,
+                        TfLiteFusedActivation activation, size_t groups,
+                        size_t group_input_channels,
+                        size_t group_output_channels,
+                        size_t output_tensor_index) {
+  bool is_quantized = IsQuantized(input_tensor);
+  uint32_t a_scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t a_zp_id = YNN_INVALID_VALUE_ID;
+  uint32_t b_scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t b_zp_id = YNN_INVALID_VALUE_ID;
+  uint32_t out_scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t out_zp_id = YNN_INVALID_VALUE_ID;
+
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+      context, subgraph, input_tensor, &a_scale_id, &a_zp_id));
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+      context, subgraph, filter_tensor, &b_scale_id, &b_zp_id));
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+      context, subgraph, output_tensor, &out_scale_id, &out_zp_id));
+  int kernel_height = filter_tensor.dims->data[1];
+  int kernel_width = filter_tensor.dims->data[2];
+
+  bool is_per_channel = false;
+  if (is_quantized) {
+    const auto* quant_params = static_cast<const TfLiteAffineQuantization*>(
+        filter_tensor.quantization.params);
+    is_per_channel =
+        quant_params && quant_params->scale && quant_params->scale->size > 1;
+  }
+
+  float padding_value = 0.0f;
+  if (is_quantized) {
+    const auto* quant_params = static_cast<const TfLiteAffineQuantization*>(
+        input_tensor.quantization.params);
+    if (quant_params && quant_params->zero_point) {
+      padding_value = quant_params->zero_point->data[0];
+    } else {
+      padding_value = input_tensor.params.zero_point;
+    }
+  }
+
+  uint32_t stencil_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DefineYnnStencil(
+      context, subgraph, input_tensor, input_id, kernel_height, kernel_width,
+      stride_height, stride_width, dilation_height, dilation_width, padding,
+      padding_value, &stencil_id));
+
+  uint32_t current_input_id = stencil_id;
+  uint32_t current_filter_id = filter_id;
+  uint32_t current_bias_id = bias_id;
+  uint32_t current_b_scale_id = b_scale_id;
+  uint32_t current_b_zp_id = b_zp_id;
+
+  if (groups != 1) {
+    // Split input: [n, h, w, kh, kw, ci] -> [n, h, w, kh, kw, g, 1, ci/g]
+    uint32_t split_input_id = YNN_INVALID_VALUE_ID;
+    const size_t input_split[] = {groups, 1, group_input_channels};
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_split_dim(
+        subgraph, 5, 3, input_split, current_input_id, &split_input_id, 0));
+    current_input_id = split_input_id;
+
+    // Split filter: [co, kh, kw, ci/g] -> [g, co/g, kh, kw, ci/g]
+    uint32_t split_filter_id = YNN_INVALID_VALUE_ID;
+    const size_t filter_split[] = {groups, group_output_channels};
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_split_dim(
+        subgraph, 0, 2, filter_split, current_filter_id, &split_filter_id, 0));
+    current_filter_id = split_filter_id;
+
+    // Split bias if present: [co] -> [g, 1, co/g]
+    if (current_bias_id != YNN_INVALID_VALUE_ID) {
+      uint32_t split_bias_id = YNN_INVALID_VALUE_ID;
+      const size_t bias_split[] = {groups, 1, group_output_channels};
+      TF_LITE_ENSURE_YNN_STATUS(ynn_define_split_dim(
+          subgraph, 0, 3, bias_split, current_bias_id, &split_bias_id, 0));
+      current_bias_id = split_bias_id;
+    }
+
+    if (is_quantized && is_per_channel) {
+      if (current_b_scale_id != YNN_INVALID_VALUE_ID) {
+        uint32_t split_scale_id = YNN_INVALID_VALUE_ID;
+        const size_t scale_split[] = {groups, group_output_channels};
+        TF_LITE_ENSURE_YNN_STATUS(
+            ynn_define_split_dim(subgraph, 0, 2, scale_split,
+                                 current_b_scale_id, &split_scale_id, 0));
+        current_b_scale_id = split_scale_id;
+      }
+      if (current_b_zp_id != YNN_INVALID_VALUE_ID) {
+        uint32_t split_zp_id = YNN_INVALID_VALUE_ID;
+        const size_t zp_split[] = {groups, group_output_channels};
+        TF_LITE_ENSURE_YNN_STATUS(ynn_define_split_dim(
+            subgraph, 0, 2, zp_split, current_b_zp_id, &split_zp_id, 0));
+        current_b_zp_id = split_zp_id;
+      }
+    }
+
+    // Transpose input: [n, h, w, kh, kw, g, 1, ci/g] -> [n, h, w, g, 1, kh, kw,
+    // ci/g]
+    uint32_t transposed_input_id = YNN_INVALID_VALUE_ID;
+    const int32_t input_perm[] = {0, 1, 2, 5, 6, 3, 4, 7};
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose(
+        subgraph, 8, input_perm, current_input_id, &transposed_input_id, 0));
+    current_input_id = transposed_input_id;
+  }
+
+  // Transpose filter:
+  // If groups == 1: [co, kh, kw, ci] -> [kh, kw, ci, co]
+  // If groups > 1: [g, co/g, kh, kw, ci/g] -> [g, kh, kw, ci/g, co/g]
+  uint32_t transposed_filter_id = YNN_INVALID_VALUE_ID;
+  if (groups == 1) {
+    int32_t swap_co_ci[4] = {1, 2, 3, 0};
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose(
+        subgraph, 4, swap_co_ci, current_filter_id, &transposed_filter_id, 0));
+  } else {
+    int32_t swap_co_ci[5] = {0, 2, 3, 4, 1};
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose(
+        subgraph, 5, swap_co_ci, current_filter_id, &transposed_filter_id, 0));
+  }
+
+  uint32_t conv_output_id = output_id;
+  if (activation != kTfLiteActNone) {
+    conv_output_id = YNN_INVALID_VALUE_ID;
+  }
+
+  uint32_t dot_output_id =
+      (groups == 1) ? conv_output_id : YNN_INVALID_VALUE_ID;
+
+  if (is_quantized) {
+    int32_t a_reduce_axes[3];
+    int32_t b_reduce_axes[3];
+    std::iota(a_reduce_axes, a_reduce_axes + 3, groups == 1 ? 3 : 5);
+    std::iota(b_reduce_axes, b_reduce_axes + 3, groups == 1 ? 0 : 1);
+
+    TF_LITE_ENSURE_STATUS(DefineQuantizedDot(
+        context, subgraph,
+        /*rank_a=*/(groups == 1) ? 6 : 8,
+        /*rank_b=*/(groups == 1) ? 4 : 5, a_reduce_axes, b_reduce_axes,
+        current_input_id, a_scale_id, a_zp_id, transposed_filter_id,
+        current_b_scale_id, current_b_zp_id, current_bias_id,
+        (groups == 1) ? 1 : 3, out_scale_id, out_zp_id, is_per_channel,
+        /*is_conv=*/true, GetYnnType(output_tensor.type), &dot_output_id));
+  } else {
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_dot(subgraph, 3, current_input_id, transposed_filter_id,
+                       current_bias_id, &dot_output_id, 0));
+  }
+
+  if (groups != 1) {
+    // Fuse [n, h, w, g, 1, co/g] -> [n, h, w, co]
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_fuse_dim(subgraph, 3, 3, dot_output_id, &conv_output_id, 0));
+  } else {
+    conv_output_id = dot_output_id;
+  }
+
+  if (activation != kTfLiteActNone) {
+    TF_LITE_ENSURE_STATUS(ApplyActivation(
+        context, subgraph, activation, conv_output_id, output_id,
+        output_tensor_index, GetYnnType(output_tensor.type)));
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace
+
+TfLiteStatus DefineConvNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                            TensorToValueIdMap& tensor_to_value_id,
+                            const NodeInfo& node) {
+  int input_tensor_index = node.inputs[0];
+  int filter_tensor_index = node.inputs[1];
+  int bias_tensor_index = node.inputs[2];
+  int output_tensor_index = node.outputs[0];
+
+  const TfLiteTensor& input_tensor = context->tensors[input_tensor_index];
+  const TfLiteTensor& filter_tensor = context->tensors[filter_tensor_index];
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+
+  uint32_t input_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                         input_tensor_index);
+  uint32_t filter_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                          filter_tensor_index);
+  uint32_t bias_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                        bias_tensor_index);
+  uint32_t output_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                          output_tensor_index);
+
+  TF_LITE_ENSURE(context, input_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, filter_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, bias_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, output_id != YNN_INVALID_VALUE_ID);
+
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLiteConvParams*>(tflite_node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  int output_channels = filter_tensor.dims->data[0];
+  int input_channels_per_group = filter_tensor.dims->data[3];
+  int input_channels = input_tensor.dims->data[3];
+  int groups = input_channels / input_channels_per_group;
+  int group_input_channels = input_channels_per_group;
+  int group_output_channels = output_channels / groups;
+
+  TF_LITE_ENSURE_STATUS(DefineConv(
+      context, subgraph, input_id, filter_id, bias_id, output_id, input_tensor,
+      filter_tensor, output_tensor, params->stride_height, params->stride_width,
+      params->dilation_height_factor, params->dilation_width_factor,
+      params->padding, params->activation, groups, group_input_channels,
+      group_output_channels, output_tensor_index));
+
+  tensor_to_value_id[output_tensor_index] = output_id;
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineDepthwiseConvNode(TfLiteContext* context,
+                                     ynn_subgraph_t subgraph,
+                                     TensorToValueIdMap& tensor_to_value_id,
+                                     const NodeInfo& node) {
+  int input_tensor_index = node.inputs[0];
+  int filter_tensor_index = node.inputs[1];
+  int bias_tensor_index = node.inputs[2];
+  int output_tensor_index = node.outputs[0];
+
+  const TfLiteTensor& input_tensor = context->tensors[input_tensor_index];
+  const TfLiteTensor& filter_tensor = context->tensors[filter_tensor_index];
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+
+  uint32_t input_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                         input_tensor_index);
+  uint32_t filter_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                          filter_tensor_index);
+  uint32_t bias_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                        bias_tensor_index);
+  uint32_t output_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                          output_tensor_index);
+
+  TF_LITE_ENSURE(context, input_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, filter_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, bias_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, output_id != YNN_INVALID_VALUE_ID);
+
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLiteDepthwiseConvParams*>(tflite_node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  int filter_channels = filter_tensor.dims->data[3];
+  int input_channels = input_tensor.dims->data[3];
+  int depth_multiplier = filter_channels / input_channels;
+
+  // Transpose filter: [1, kh, kw, ci * dm] -> [ci * dm, kh, kw, 1]
+  const int32_t swap_dims[4] = {3, 1, 2, 0};
+  uint32_t transposed_filter_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose(
+      subgraph, 4, swap_dims, filter_id, &transposed_filter_id, 0));
+
+  int groups = input_channels;
+  int group_input_channels = 1;
+  int group_output_channels = depth_multiplier;
+  TF_LITE_ENSURE_STATUS(DefineConv(
+      context, subgraph, input_id, transposed_filter_id, bias_id, output_id,
+      input_tensor, filter_tensor, output_tensor, params->stride_height,
+      params->stride_width, params->dilation_height_factor,
+      params->dilation_width_factor, params->padding, params->activation,
+      groups, group_input_channels, group_output_channels,
+      output_tensor_index));
+
+  tensor_to_value_id[output_tensor_index] = output_id;
+  return kTfLiteOk;
+}
+
+}  // namespace ynnpack
+}  // namespace tflite

tensorflow/lite/delegates/ynnpack/dot.h

@@ -0,0 +1,63 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_DOT_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_DOT_H_
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsBatchMatMulSupported(const TfLiteRegistration* registration,
+                                    const TfLiteNode* node,
+                                    TfLiteContext* context);
+
+TfLiteStatus IsFullyConnectedSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context);
+
+TfLiteStatus DefineBatchMatMulNode(TfLiteContext* context,
+                                   ynn_subgraph_t subgraph,
+                                   TensorToValueIdMap& tensor_to_value_id,
+                                   const NodeInfo& node);
+
+TfLiteStatus DefineFullyConnectedNode(TfLiteContext* context,
+                                      ynn_subgraph_t subgraph,
+                                      TensorToValueIdMap& tensor_to_value_id,
+                                      const NodeInfo& node);
+
+TfLiteStatus IsConvSupported(const TfLiteRegistration* registration,
+                             const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus IsDepthwiseConvSupported(const TfLiteRegistration* registration,
+                                      const TfLiteNode* node,
+                                      TfLiteContext* context);
+
+TfLiteStatus DefineConvNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                            TensorToValueIdMap& tensor_to_value_id,
+                            const NodeInfo& node);
+
+TfLiteStatus DefineDepthwiseConvNode(TfLiteContext* context,
+                                     ynn_subgraph_t subgraph,
+                                     TensorToValueIdMap& tensor_to_value_id,
+                                     const NodeInfo& node);
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_DOT_H_

tensorflow/lite/delegates/ynnpack/elementwise.cc

@@ -0,0 +1,492 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/ynnpack/elementwise.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "ynnpack/composites/composites.h"  // from @XNNPACK
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+// ==============================================================================
+// Unary Operations
+// ==============================================================================
+
+TfLiteStatus IsUnaryOpSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node,
+                                TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 1);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+  TF_LITE_ENSURE(context, YnnTypeElementCount(GetYnnType(input.type)) == 1);
+  TF_LITE_ENSURE(context, YnnTypeElementCount(GetYnnType(output.type)) == 1);
+
+  int builtin_code = registration->builtin_code;
+  ynn_unary_operator op = GetYnnUnaryOperator(builtin_code);
+  bool is_decomposed = (builtin_code == kTfLiteBuiltinGelu ||
+                        builtin_code == kTfLiteBuiltinElu ||
+                        builtin_code == kTfLiteBuiltinLeakyRelu ||
+                        builtin_code == kTfLiteBuiltinHardSwish ||
+                        builtin_code == kTfLiteBuiltinRelu ||
+                        builtin_code == kTfLiteBuiltinRelu6 ||
+                        builtin_code == kTfLiteBuiltinReluN1To1 ||
+                        builtin_code == kTfLiteBuiltinRelu0To1);
+  TF_LITE_ENSURE(context, op != ynn_unary_invalid || is_decomposed);
+
+  if (op != ynn_unary_convert) {
+    TF_LITE_ENSURE_EQ(context, input.type, output.type);
+  }
+
+  if (op == ynn_unary_convert) {
+    // Reject constant inputs to allow TFLite caching optimization.
+    TF_LITE_ENSURE_MSG(context, input.allocation_type != kTfLiteMmapRo,
+                       "Constant input for convert is not supported");
+    // YNNPACK convert to integer uses rounding, but TFLite Cast expects
+    // truncation. Reject all float-to-integer conversions.
+    if ((input.type == kTfLiteFloat32 || input.type == kTfLiteFloat16) &&
+        (output.type == kTfLiteInt32 || output.type == kTfLiteInt8 ||
+         output.type == kTfLiteUInt8 || output.type == kTfLiteInt16)) {
+      TF_LITE_ENSURE_MSG(context, false,
+                         "Float to integer conversion is not supported");
+    }
+  }
+
+  if (op == ynn_unary_square_root || op == ynn_unary_reciprocal_square_root) {
+    // Reject quantized Sqrt/Rsqrt because we cannot report errors for negative
+    // inputs during execution if we delegate them.
+    TF_LITE_ENSURE_MSG(context,
+                       input.quantization.type == kTfLiteNoQuantization,
+                       "Quantized Sqrt/Rsqrt is not supported");
+  }
+
+  // Check for fused activation.
+  TfLiteFusedActivation activation = GetFusedActivation(registration, node);
+  TF_LITE_ENSURE(context, IsActivationSupported(activation, output.type));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineUnaryNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                             TensorToValueIdMap& tensor_to_value_id,
+                             const NodeInfo& node) {
+  return DefineDecomposedUnaryNode(
+      context, subgraph, tensor_to_value_id, node,
+      [context, subgraph, &node](uint32_t input_id,
+                                 uint32_t& output_id) -> TfLiteStatus {
+        ynn_unary_operator op = GetYnnUnaryOperator(node.builtin_code);
+        if (op != ynn_unary_invalid) {
+          TF_LITE_ENSURE_YNN_STATUS(
+              ynn_define_unary(subgraph, op, input_id, &output_id, 0));
+          return kTfLiteOk;
+        }
+
+        switch (node.builtin_code) {
+          case kTfLiteBuiltinGelu: {
+            TfLiteNode* tflite_node;
+            TfLiteRegistration* reg;
+            TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+                context, node.node_index, &tflite_node, &reg));
+            const auto* params =
+                static_cast<const TfLiteGeluParams*>(tflite_node->builtin_data);
+            bool approximate = params && params->approximate;
+            if (approximate) {
+              TF_LITE_ENSURE_YNN_STATUS(
+                  ynn::define_approx_gelu(subgraph, input_id, output_id));
+            } else {
+              TF_LITE_ENSURE_YNN_STATUS(
+                  ynn::define_gelu(subgraph, input_id, output_id));
+            }
+            return kTfLiteOk;
+          }
+          case kTfLiteBuiltinElu:
+            TF_LITE_ENSURE_YNN_STATUS(
+                ynn::define_elu(subgraph, input_id, 1.0f, output_id));
+            return kTfLiteOk;
+          case kTfLiteBuiltinLeakyRelu: {
+            TfLiteNode* tflite_node;
+            TfLiteRegistration* reg;
+            TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+                context, node.node_index, &tflite_node, &reg));
+            const auto* params = static_cast<const TfLiteLeakyReluParams*>(
+                tflite_node->builtin_data);
+            float alpha = params ? params->alpha : 0.2f;
+            TF_LITE_ENSURE_YNN_STATUS(
+                ynn::define_leaky_relu(subgraph, input_id, alpha, output_id));
+            return kTfLiteOk;
+          }
+          case kTfLiteBuiltinHardSwish:
+            TF_LITE_ENSURE_YNN_STATUS(
+                ynn::define_hardswish(subgraph, input_id, output_id));
+            return kTfLiteOk;
+          case kTfLiteBuiltinRelu0To1:
+            return ApplyClamp(context, subgraph, 0.0, 1.0, input_id, output_id,
+                              node.outputs[0], ynn_type_fp32);
+          case kTfLiteBuiltinRelu:
+          case kTfLiteBuiltinRelu6:
+          case kTfLiteBuiltinReluN1To1: {
+            TfLiteFusedActivation activation = kTfLiteActNone;
+            if (node.builtin_code == kTfLiteBuiltinRelu) {
+              activation = kTfLiteActRelu;
+            } else if (node.builtin_code == kTfLiteBuiltinRelu6) {
+              activation = kTfLiteActRelu6;
+            } else if (node.builtin_code == kTfLiteBuiltinReluN1To1) {
+              activation = kTfLiteActReluN1To1;
+            }
+            return ApplyActivation(context, subgraph, activation, input_id,
+                                   output_id, node.outputs[0], ynn_type_fp32);
+          }
+          default:
+            TF_LITE_ENSURE(context, false);
+        }
+      });
+}
+
+// ==============================================================================
+// Binary Operations
+// ==============================================================================
+
+TfLiteStatus IsBinaryOpSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 2);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input1 = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& input2 = context->tensors[node->inputs->data[1]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input1));
+  TF_LITE_ENSURE(context, IsTensorSupported(input2));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+  TF_LITE_ENSURE(context, YnnTypeElementCount(GetYnnType(input1.type)) == 1);
+  TF_LITE_ENSURE(context, YnnTypeElementCount(GetYnnType(input2.type)) == 1);
+  TF_LITE_ENSURE(context, YnnTypeElementCount(GetYnnType(output.type)) == 1);
+
+  // YNNPACK division is floating point division only. Integer and quantized
+  // division are not supported.
+  if ((registration->builtin_code == kTfLiteBuiltinDiv ||
+       registration->builtin_code == kTfLiteBuiltinStablehloDivide) &&
+      output.type != kTfLiteFloat32 && output.type != kTfLiteFloat16 &&
+      output.type != kTfLiteBFloat16) {
+    TF_LITE_ENSURE_MSG(
+        context, false,
+        "Non-floating point division is not supported by YNNPACK");
+  }
+
+  // Check for fused activation.
+  TfLiteFusedActivation activation = GetFusedActivation(registration, node);
+  TF_LITE_ENSURE(context, IsActivationSupported(activation, output.type));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineBinaryNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                              TensorToValueIdMap& tensor_to_value_id,
+                              const NodeInfo& node) {
+  TF_LITE_ENSURE_EQ(context, node.inputs.size(), 2);
+  TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1);
+  int input1_tensor_index = node.inputs[0];
+  int input2_tensor_index = node.inputs[1];
+  int output_tensor_index = node.outputs[0];
+
+  uint32_t input1_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input1_tensor_index);
+  uint32_t input2_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input2_tensor_index);
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+
+  TF_LITE_ENSURE(context, input1_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, input2_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, output_val_id != YNN_INVALID_VALUE_ID);
+
+  uint32_t float_input1_val_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DequantizeIfNeeded(
+      context, subgraph, tensor_to_value_id, input1_tensor_index, input1_val_id,
+      &float_input1_val_id));
+
+  uint32_t float_input2_val_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DequantizeIfNeeded(
+      context, subgraph, tensor_to_value_id, input2_tensor_index, input2_val_id,
+      &float_input2_val_id));
+
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+  ynn_type internal_type = IsQuantized(output_tensor)
+                               ? ynn_type_fp32
+                               : GetYnnType(output_tensor.type);
+  TfLiteFusedActivation activation = node.activation;
+  bool is_output_quantized = IsQuantized(output_tensor);
+
+  uint32_t float_output_val_id =
+      !is_output_quantized && activation == kTfLiteActNone
+          ? output_val_id
+          : YNN_INVALID_VALUE_ID;
+
+  uint32_t broadcasted_input1_val_id = float_input1_val_id;
+  uint32_t broadcasted_input2_val_id = float_input2_val_id;
+
+  const TfLiteTensor& input1_tensor = context->tensors[input1_tensor_index];
+  const TfLiteTensor& input2_tensor = context->tensors[input2_tensor_index];
+  int rank1 = input1_tensor.dims->size;
+  int rank2 = input2_tensor.dims->size;
+
+  if (!IsStablehloOp(node.builtin_code)) {
+    // If the op is not a StableHLO op, we might need to broadcast.
+    TF_LITE_ENSURE_STATUS(ImplementMutualBroadcasting(
+        context, subgraph, rank1, rank2, 0, 0, broadcasted_input1_val_id,
+        broadcasted_input2_val_id));
+  }
+
+  TF_LITE_ENSURE_YNN_STATUS(
+      ynn_define_binary(subgraph, GetYnnBinaryOperator(node.builtin_code),
+                        broadcasted_input1_val_id, broadcasted_input2_val_id,
+                        &float_output_val_id, 0));
+
+  uint32_t active_output_val_id = float_output_val_id;
+  if (activation != kTfLiteActNone) {
+    uint32_t activation_output_val_id =
+        is_output_quantized ? YNN_INVALID_VALUE_ID : output_val_id;
+    TF_LITE_ENSURE_STATUS(ApplyActivation(
+        context, subgraph, activation, float_output_val_id,
+        activation_output_val_id, output_tensor_index, internal_type));
+    active_output_val_id = activation_output_val_id;
+  }
+
+  if (is_output_quantized) {
+    TF_LITE_ENSURE_STATUS(Quantize(context, subgraph, tensor_to_value_id,
+                                   output_tensor_index, active_output_val_id,
+                                   output_val_id));
+  }
+
+  return kTfLiteOk;
+}
+
+// ==============================================================================
+// Ternary Operations (Clamp)
+// ==============================================================================
+
+TfLiteStatus IsStablehloClampSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 3);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& min_tensor = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& operand_tensor = context->tensors[node->inputs->data[1]];
+  const TfLiteTensor& max_tensor = context->tensors[node->inputs->data[2]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(min_tensor));
+  TF_LITE_ENSURE(context, IsTensorSupported(operand_tensor));
+  TF_LITE_ENSURE(context, IsTensorSupported(max_tensor));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+
+  ynn_type operand_ynn_type = GetYnnType(operand_tensor.type);
+  switch (operand_ynn_type) {
+    case ynn_type_int8:
+    case ynn_type_uint8:
+    case ynn_type_int32:
+    case ynn_type_fp16:
+    case ynn_type_bf16:
+      return kTfLiteError;
+    default:
+      break;
+  }
+
+  TF_LITE_ENSURE_EQ(context, min_tensor.type, output.type);
+  TF_LITE_ENSURE_EQ(context, operand_tensor.type, output.type);
+  TF_LITE_ENSURE_EQ(context, max_tensor.type, output.type);
+
+  // Fused activation is not supported for StableHLO Clamp.
+  TfLiteFusedActivation activation = GetFusedActivation(registration, node);
+  TF_LITE_ENSURE_EQ(context, activation, kTfLiteActNone);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineStablehloClampNode(TfLiteContext* context,
+                                      ynn_subgraph_t subgraph,
+                                      TensorToValueIdMap& tensor_to_value_id,
+                                      const NodeInfo& node) {
+  TF_LITE_ENSURE_EQ(context, node.inputs.size(), 3);
+  TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1);
+
+  int min_tensor_index = node.inputs[0];
+  int operand_tensor_index = node.inputs[1];
+  int max_tensor_index = node.inputs[2];
+  int output_tensor_index = node.outputs[0];
+
+  uint32_t min_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, min_tensor_index);
+  uint32_t operand_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, operand_tensor_index);
+  uint32_t max_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, max_tensor_index);
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+
+  TF_LITE_ENSURE(context, min_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, operand_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, max_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, output_val_id != YNN_INVALID_VALUE_ID);
+
+  uint32_t float_min_val_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DequantizeIfNeeded(context, subgraph,
+                                           tensor_to_value_id, min_tensor_index,
+                                           min_val_id, &float_min_val_id));
+
+  uint32_t float_operand_val_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DequantizeIfNeeded(
+      context, subgraph, tensor_to_value_id, operand_tensor_index,
+      operand_val_id, &float_operand_val_id));
+
+  uint32_t float_max_val_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DequantizeIfNeeded(context, subgraph,
+                                           tensor_to_value_id, max_tensor_index,
+                                           max_val_id, &float_max_val_id));
+
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+  bool is_output_quantized = IsQuantized(output_tensor);
+
+  uint32_t float_output_val_id =
+      !is_output_quantized ? output_val_id : YNN_INVALID_VALUE_ID;
+
+  // temp = max(operand, min)
+  uint32_t temp_val_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(
+      ynn_define_binary(subgraph, ynn_binary_max, float_operand_val_id,
+                        float_min_val_id, &temp_val_id, 0));
+
+  // output = min(temp, max)
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_binary(subgraph, ynn_binary_min,
+                                              temp_val_id, float_max_val_id,
+                                              &float_output_val_id, 0));
+
+  if (is_output_quantized) {
+    TF_LITE_ENSURE_STATUS(Quantize(context, subgraph, tensor_to_value_id,
+                                   output_tensor_index, float_output_val_id,
+                                   output_val_id));
+  }
+
+  return kTfLiteOk;
+}
+
+// ==============================================================================
+// Quantize/Dequantize Operations
+// ==============================================================================
+
+TfLiteStatus IsQuantizeSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 1);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+  TF_LITE_ENSURE(context, !IsQuantized(input));
+  TF_LITE_ENSURE(context, YnnTypeElementCount(GetYnnType(output.type)) == 1);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineQuantizeNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                                TensorToValueIdMap& tensor_to_value_id,
+                                const NodeInfo& node) {
+  TF_LITE_ENSURE_EQ(context, node.inputs.size(), 1);
+  TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1);
+
+  int input_tensor_index = node.inputs[0];
+  int output_tensor_index = node.outputs[0];
+
+  uint32_t input_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input_tensor_index);
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+  uint32_t scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t zp_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+      context, subgraph, output_tensor, &scale_id, &zp_id));
+
+  ynn_type ynn_type = GetYnnType(output_tensor.type);
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_quantize(
+      subgraph, input_val_id, ynn_type, zp_id, scale_id, &output_val_id, 0));
+  return kTfLiteOk;
+}
+
+TfLiteStatus IsDequantizeSupported(const TfLiteRegistration* registration,
+                                   const TfLiteNode* node,
+                                   TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 1);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+  TF_LITE_ENSURE(context, !IsQuantized(output));
+  TF_LITE_ENSURE(context, YnnTypeElementCount(GetYnnType(input.type)) == 1);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineDequantizeNode(TfLiteContext* context,
+                                  ynn_subgraph_t subgraph,
+                                  TensorToValueIdMap& tensor_to_value_id,
+                                  const NodeInfo& node) {
+  TF_LITE_ENSURE_EQ(context, node.inputs.size(), 1);
+  TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1);
+
+  int input_tensor_index = node.inputs[0];
+  int output_tensor_index = node.outputs[0];
+
+  uint32_t input_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input_tensor_index);
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+
+  const TfLiteTensor& input_tensor = context->tensors[input_tensor_index];
+
+  uint32_t scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t zp_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+      context, subgraph, input_tensor, &scale_id, &zp_id));
+
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_dequantize(subgraph, input_val_id, zp_id,
+                                                  scale_id, ynn_type_fp32,
+                                                  &output_val_id, 0));
+  return kTfLiteOk;
+}
+
+}  // namespace ynnpack
+}  // namespace tflite

tensorflow/lite/delegates/ynnpack/elementwise.h

@@ -0,0 +1,72 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_ELEMENTWISE_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_ELEMENTWISE_H_
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+// Unary operations
+TfLiteStatus IsUnaryOpSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus DefineUnaryNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                             TensorToValueIdMap& tensor_to_value_id,
+                             const NodeInfo& node);
+
+// Binary operations
+TfLiteStatus IsBinaryOpSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context);
+
+TfLiteStatus DefineBinaryNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                              TensorToValueIdMap& tensor_to_value_id,
+                              const NodeInfo& node);
+
+// Ternary operations (Clamp)
+TfLiteStatus IsStablehloClampSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context);
+
+TfLiteStatus DefineStablehloClampNode(TfLiteContext* context,
+                                      ynn_subgraph_t subgraph,
+                                      TensorToValueIdMap& tensor_to_value_id,
+                                      const NodeInfo& node);
+
+// Quantize/Dequantize operations
+TfLiteStatus IsQuantizeSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context);
+TfLiteStatus DefineQuantizeNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                                TensorToValueIdMap& tensor_to_value_id,
+                                const NodeInfo& node);
+
+TfLiteStatus IsDequantizeSupported(const TfLiteRegistration* registration,
+                                   const TfLiteNode* node,
+                                   TfLiteContext* context);
+TfLiteStatus DefineDequantizeNode(TfLiteContext* context,
+                                  ynn_subgraph_t subgraph,
+                                  TensorToValueIdMap& tensor_to_value_id,
+                                  const NodeInfo& node);
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_ELEMENTWISE_H_

tensorflow/lite/delegates/ynnpack/pooling.cc

@@ -0,0 +1,151 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/ynnpack/pooling.h"
+
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#include "ynnpack/composites/composites.h"  // from @XNNPACK
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsPoolingSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node,
+                                TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 1);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+
+  TF_LITE_ENSURE_EQ(context, input.type, output.type);
+  TF_LITE_ENSURE(context, QuantizationParamsEqual(input, output));
+
+  // We only support NHWC format for 2D pooling, which means 4D tensor. NHWC
+  // format NHWC is represented as [batch, height, width, channels].
+  TF_LITE_ENSURE_EQ(context, input.dims->size, 4);
+
+  const auto* params = static_cast<const TfLitePoolParams*>(node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+  TF_LITE_ENSURE(context, params->stride_height > 0);
+  TF_LITE_ENSURE(context, params->stride_width > 0);
+  TF_LITE_ENSURE(context, params->filter_height > 0);
+  TF_LITE_ENSURE(context, params->filter_width > 0);
+
+  TF_LITE_ENSURE(context,
+                 IsActivationSupported(params->activation, output.type));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineMaxPool2DNode(TfLiteContext* context,
+                                 ynn_subgraph_t subgraph,
+                                 TensorToValueIdMap& tensor_to_value_id,
+                                 const NodeInfo& node) {
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLitePoolParams*>(tflite_node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  int input_tensor_index = node.inputs[0];
+  int output_tensor_index = node.outputs[0];
+  const TfLiteTensor& input_tensor = context->tensors[input_tensor_index];
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+
+  uint32_t input_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                         input_tensor_index);
+  uint32_t output_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
+                                          output_tensor_index);
+
+  TF_LITE_ENSURE(context, input_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, output_id != YNN_INVALID_VALUE_ID);
+
+  TfLiteFusedActivation activation = node.activation;
+
+  uint32_t maxpool_output_id =
+      activation == kTfLiteActNone ? output_id : YNN_INVALID_VALUE_ID;
+
+  // Pad with -inf for max pooling.
+  float padding_val = -std::numeric_limits<float>::infinity();
+  if (IsQuantized(input_tensor)) {
+    if (input_tensor.type == kTfLiteInt8) {
+      padding_val = std::numeric_limits<int8_t>::min();
+    } else if (input_tensor.type == kTfLiteUInt8) {
+      padding_val = std::numeric_limits<uint8_t>::min();
+    }
+  }
+
+  uint32_t stencil_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DefineYnnStencil(
+      context, subgraph, input_tensor, input_id, params->filter_height,
+      params->filter_width, params->stride_height, params->stride_width,
+      /*dilation_height=*/1, /*dilation_width=*/1, params->padding, padding_val,
+      &stencil_id));
+
+  const int32_t reduce_axes[] = {3, 4};
+  TF_LITE_ENSURE_YNN_STATUS(
+      ynn_define_reduce(subgraph, ynn_reduce_max, 2, reduce_axes, stencil_id,
+                        YNN_INVALID_VALUE_ID, &maxpool_output_id, /*flags=*/0));
+
+  if (activation != kTfLiteActNone) {
+    TF_LITE_ENSURE_STATUS(ApplyActivation(
+        context, subgraph, activation, maxpool_output_id, output_id,
+        output_tensor_index, GetYnnType(output_tensor.type)));
+  }
+
+  tensor_to_value_id[output_tensor_index] = output_id;
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineAveragePool2DNode(TfLiteContext* context,
+                                     ynn_subgraph_t subgraph,
+                                     TensorToValueIdMap& tensor_to_value_id,
+                                     const NodeInfo& node) {
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLitePoolParams*>(tflite_node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  return DefineDecomposedUnaryNode(
+      context, subgraph, tensor_to_value_id, node,
+      [context, subgraph, params](uint32_t input_id,
+                                  uint32_t& output_id) -> TfLiteStatus {
+        bool padding_same = (params->padding == kTfLitePaddingSame);
+        TF_LITE_ENSURE_YNN_STATUS(ynn::define_average_pool_2d(
+            subgraph, input_id, ynn_type_fp32, padding_same,
+            params->filter_height, params->filter_width, params->stride_height,
+            params->stride_width, output_id));
+        return kTfLiteOk;
+      });
+}
+
+}  // namespace ynnpack
+}  // namespace tflite

tensorflow/lite/delegates/ynnpack/pooling.h

@@ -0,0 +1,42 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_POOLING_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_POOLING_H_
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsPoolingSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus DefineMaxPool2DNode(TfLiteContext* context,
+                                 ynn_subgraph_t subgraph,
+                                 TensorToValueIdMap& tensor_to_value_id,
+                                 const NodeInfo& node);
+
+TfLiteStatus DefineAveragePool2DNode(TfLiteContext* context,
+                                     ynn_subgraph_t subgraph,
+                                     TensorToValueIdMap& tensor_to_value_id,
+                                     const NodeInfo& node);
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_POOLING_H_

tensorflow/lite/delegates/ynnpack/reduction.cc

@@ -0,0 +1,134 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/ynnpack/reduction.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "ynnpack/composites/composites.h"  // from @XNNPACK
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsReductionSupported(const TfLiteRegistration* registration,
+                                  const TfLiteNode* node,
+                                  TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 2);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  TF_LITE_ENSURE(context, NumElements(&input) > 0);
+  const TfLiteTensor& axes = context->tensors[node->inputs->data[1]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+
+  TF_LITE_ENSURE_EQ(context, input.type, output.type);
+
+  TF_LITE_ENSURE(context, QuantizationParamsEqual(input, output));
+
+  TF_LITE_ENSURE(context, axes.allocation_type == kTfLiteMmapRo);
+  TF_LITE_ENSURE_EQ(context, axes.type, kTfLiteInt32);
+
+  // YNNPACK reduce only supports 0D or 1D axes for now.
+  TF_LITE_ENSURE(context, axes.dims->size <= 1);
+  TF_LITE_ENSURE_MSG(context, input.dims->size <= YNN_MAX_TENSOR_RANK,
+                     "Input rank exceeds max rank");
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineReductionNode(TfLiteContext* context,
+                                 ynn_subgraph_t subgraph,
+                                 TensorToValueIdMap& tensor_to_value_id,
+                                 const NodeInfo& node) {
+  TF_LITE_ENSURE_EQ(context, node.inputs.size(), 2);
+  TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1);
+
+  int input_tensor_index = node.inputs[0];
+  int axes_tensor_index = node.inputs[1];
+  int output_tensor_index = node.outputs[0];
+
+  const TfLiteTensor& input_tensor = context->tensors[input_tensor_index];
+  const TfLiteTensor& axes_tensor = context->tensors[axes_tensor_index];
+
+  uint32_t input_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input_tensor_index);
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+
+  TF_LITE_ENSURE(context, input_val_id != YNN_INVALID_VALUE_ID);
+  TF_LITE_ENSURE(context, output_val_id != YNN_INVALID_VALUE_ID);
+
+  int num_axes = axes_tensor.dims->size == 0 ? 1 : axes_tensor.dims->data[0];
+  const int32_t* axes_data =
+      reinterpret_cast<const int32_t*>(axes_tensor.data.raw);
+
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLiteReducerParams*>(tflite_node->builtin_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+
+  bool is_quantized = IsQuantized(input_tensor);
+  bool is_mean = (node.builtin_code == kTfLiteBuiltinMean);
+  bool is_sum = (node.builtin_code == kTfLiteBuiltinSum);
+
+  if (is_mean || is_sum) {
+    const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+    ynn_type output_type = GetYnnType(output_tensor.type);
+
+    uint32_t input_scale_id = YNN_INVALID_VALUE_ID;
+    uint32_t input_zp_id = YNN_INVALID_VALUE_ID;
+    if (is_quantized) {
+      TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+          context, subgraph, input_tensor, &input_scale_id, &input_zp_id));
+    }
+
+    uint32_t output_scale_id = YNN_INVALID_VALUE_ID;
+    uint32_t output_zp_id = YNN_INVALID_VALUE_ID;
+    if (is_quantized) {
+      TF_LITE_ENSURE_STATUS(DefineQuantizationParams(
+          context, subgraph, output_tensor, &output_scale_id, &output_zp_id));
+    }
+
+    TF_LITE_ENSURE_YNN_STATUS(ynn::define_reduce_sum(
+        subgraph, num_axes, axes_data, input_val_id, input_zp_id,
+        input_scale_id, params->keep_dims, is_mean, /*squared=*/false,
+        output_type, output_zp_id, output_scale_id, output_val_id));
+  } else {
+    // For Min/Max, we can reduce directly (same for float and quantized).
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_reduce(
+        subgraph, GetYnnReduceOperator(node.builtin_code), num_axes, axes_data,
+        input_val_id, YNN_INVALID_VALUE_ID, &output_val_id,
+        params->keep_dims ? YNN_NODE_FLAG_KEEP_DIMS : 0));
+  }
+
+  tensor_to_value_id[output_tensor_index] = output_val_id;
+  return kTfLiteOk;
+}
+
+}  // namespace ynnpack
+}  // namespace tflite

tensorflow/lite/delegates/ynnpack/reduction.h

@@ -0,0 +1,38 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_REDUCTION_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_REDUCTION_H_
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsReductionSupported(const TfLiteRegistration* registration,
+                                  const TfLiteNode* node,
+                                  TfLiteContext* context);
+
+TfLiteStatus DefineReductionNode(TfLiteContext* context,
+                                 ynn_subgraph_t subgraph,
+                                 TensorToValueIdMap& tensor_to_value_id,
+                                 const NodeInfo& node);
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_REDUCTION_H_

tensorflow/lite/delegates/ynnpack/softmax.cc

@@ -0,0 +1,82 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/ynnpack/softmax.h"
+
+#include <cstdint>
+
+#include "ynnpack/composites/composites.h"  // from @XNNPACK
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsSoftmaxSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node,
+                                TfLiteContext* context) {
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 1);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor& input = context->tensors[node->inputs->data[0]];
+  const TfLiteTensor& output = context->tensors[node->outputs->data[0]];
+
+  TF_LITE_ENSURE(context, IsTensorSupported(input));
+  TF_LITE_ENSURE(context, IsTensorSupported(output));
+
+  TF_LITE_ENSURE_EQ(context, input.type, output.type);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineSoftmaxNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                               TensorToValueIdMap& tensor_to_value_id,
+                               const NodeInfo& node) {
+  TfLiteNode* tflite_node;
+  TfLiteRegistration* reg;
+  TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
+      context, node.node_index, &tflite_node, &reg));
+  const auto* params =
+      static_cast<const TfLiteSoftmaxParams*>(tflite_node->builtin_data);
+  float beta = params ? params->beta : 1.0f;
+
+  return DefineDecomposedUnaryNode(
+      context, subgraph, tensor_to_value_id, node,
+      [context, subgraph, beta](uint32_t input_id,
+                                uint32_t& output_id) -> TfLiteStatus {
+        TF_LITE_ENSURE_YNN_STATUS(
+            ynn::define_softmax(subgraph, input_id, beta, output_id));
+        return kTfLiteOk;
+      });
+}
+
+TfLiteStatus DefineLogSoftmaxNode(TfLiteContext* context,
+                                  ynn_subgraph_t subgraph,
+                                  TensorToValueIdMap& tensor_to_value_id,
+                                  const NodeInfo& node) {
+  return DefineDecomposedUnaryNode(
+      context, subgraph, tensor_to_value_id, node,
+      [context, subgraph](uint32_t input_id,
+                          uint32_t& output_id) -> TfLiteStatus {
+        TF_LITE_ENSURE_YNN_STATUS(
+            ynn::define_log_softmax(subgraph, input_id, output_id));
+        return kTfLiteOk;
+      });
+}
+
+}  // namespace ynnpack
+}  // namespace tflite

tensorflow/lite/delegates/ynnpack/softmax.h

@@ -0,0 +1,41 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_SOFTMAX_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_SOFTMAX_H_
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+TfLiteStatus IsSoftmaxSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node, TfLiteContext* context);
+
+TfLiteStatus DefineSoftmaxNode(TfLiteContext* context, ynn_subgraph_t subgraph,
+                               TensorToValueIdMap& tensor_to_value_id,
+                               const NodeInfo& node);
+
+TfLiteStatus DefineLogSoftmaxNode(TfLiteContext* context,
+                                  ynn_subgraph_t subgraph,
+                                  TensorToValueIdMap& tensor_to_value_id,
+                                  const NodeInfo& node);
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_SOFTMAX_H_

tensorflow/lite/delegates/ynnpack/utils.cc

@@ -0,0 +1,778 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <numeric>
+#include <vector>
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/types/half.h"
+
+namespace tflite {
+namespace ynnpack {
+
+namespace {
+uint16_t FloatToBfloat16(float f) {
+  uint32_t val;
+  std::memcpy(&val, &f, sizeof(float));
+  // Round to nearest even.
+  uint32_t rounding_bias = 0x7fff + ((val >> 16) & 1);
+  val += rounding_bias;
+  return static_cast<uint16_t>(val >> 16);
+}
+
+template <typename T>
+T QuantizeValue(float value, float scale, int32_t zero_point) {
+  int32_t quantized = zero_point + std::round(value / scale);
+  return static_cast<T>(std::max<int32_t>(
+      std::numeric_limits<T>::min(),
+      std::min<int32_t>(std::numeric_limits<T>::max(), quantized)));
+}
+}  // namespace
+
+TfLiteStatus GetTfLiteTensorValueAsDouble(TfLiteContext* context,
+                                          const TfLiteTensor& tensor, int index,
+                                          double* value) {
+  TF_LITE_ENSURE(context, tensor.data.raw != nullptr);
+  int num_elements = tflite::NumElements(&tensor);
+  TF_LITE_ENSURE_MSG(context, index >= 0 && index < num_elements,
+                     "Index %d out of bounds for tensor with %d elements",
+                     index, num_elements);
+
+  switch (tensor.type) {
+    case kTfLiteFloat32:
+      *value = reinterpret_cast<const float*>(tensor.data.raw)[index];
+      break;
+    case kTfLiteInt32:
+      *value = reinterpret_cast<const int32_t*>(tensor.data.raw)[index];
+      break;
+    case kTfLiteFloat16:
+      *value = static_cast<float>(
+          reinterpret_cast<const tflite::half*>(tensor.data.raw)[index]);
+      break;
+    case kTfLiteInt8:
+      *value = reinterpret_cast<const int8_t*>(tensor.data.raw)[index];
+      break;
+    case kTfLiteUInt8:
+      *value = reinterpret_cast<const uint8_t*>(tensor.data.raw)[index];
+      break;
+    default:
+      TF_LITE_ENSURE_MSG(context, false, "Unsupported type %d", tensor.type);
+  }
+  return kTfLiteOk;
+}
+
+ynn_type GetYnnType(TfLiteType type) {
+  switch (type) {
+    case kTfLiteFloat32:
+      return ynn_type_fp32;
+    case kTfLiteInt32:
+      return ynn_type_int32;
+    case kTfLiteFloat16:
+      return ynn_type_fp16;
+    case kTfLiteFloat64:
+      return ynn_type_fp64;
+    case kTfLiteBFloat16:
+      return ynn_type_bf16;
+    case kTfLiteInt8:
+      return ynn_type_int8;
+    case kTfLiteUInt8:
+      return ynn_type_uint8;
+    case kTfLiteInt4:
+      return ynn_type_int4;
+    case kTfLiteUInt4:
+      return ynn_type_uint4;
+    case kTfLiteInt2:
+      return ynn_type_int2;
+    default:
+      return ynn_type_invalid;
+  }
+}
+
+ynn_unary_operator GetYnnUnaryOperator(int builtin_code) {
+  switch (builtin_code) {
+    case kTfLiteBuiltinAbs:
+    case kTfLiteBuiltinStablehloAbs:
+      return ynn_unary_abs;
+    case kTfLiteBuiltinCeil:
+      return ynn_unary_ceil;
+    case kTfLiteBuiltinCos:
+    case kTfLiteBuiltinStablehloCosine:
+      return ynn_unary_cosine;
+    case kTfLiteBuiltinExp:
+    case kTfLiteBuiltinStablehloExponential:
+      return ynn_unary_exp;
+    case kTfLiteBuiltinFloor:
+    case kTfLiteBuiltinStablehloFloor:
+      return ynn_unary_floor;
+    case kTfLiteBuiltinLog:
+    case kTfLiteBuiltinStablehloLog:
+      return ynn_unary_log;
+    case kTfLiteBuiltinNeg:
+    case kTfLiteBuiltinStablehloNegate:
+      return ynn_unary_negate;
+    case kTfLiteBuiltinRsqrt:
+    case kTfLiteBuiltinStablehloRsqrt:
+      return ynn_unary_reciprocal_square_root;
+    case kTfLiteBuiltinRound:
+      return ynn_unary_round;
+    case kTfLiteBuiltinLogistic:
+    case kTfLiteBuiltinStablehloLogistic:
+      return ynn_unary_sigmoid;
+    case kTfLiteBuiltinSin:
+      return ynn_unary_sine;
+    case kTfLiteBuiltinSquare:
+      return ynn_unary_square;
+    case kTfLiteBuiltinSqrt:
+      return ynn_unary_square_root;
+    case kTfLiteBuiltinTanh:
+    case kTfLiteBuiltinStablehloTanh:
+      return ynn_unary_tanh;
+    case kTfLiteBuiltinSign:
+      return ynn_unary_sign;
+    case kTfLiteBuiltinCast:
+    case kTfLiteBuiltinStablehloConvert:
+      return ynn_unary_convert;
+    default:
+      return ynn_unary_invalid;
+  }
+}
+
+ynn_binary_operator GetYnnBinaryOperator(int builtin_code) {
+  switch (builtin_code) {
+    case kTfLiteBuiltinAdd:
+    case kTfLiteBuiltinStablehloAdd:
+      return ynn_binary_add;
+    case kTfLiteBuiltinDiv:
+    case kTfLiteBuiltinStablehloDivide:
+      return ynn_binary_divide;
+    case kTfLiteBuiltinMaximum:
+    case kTfLiteBuiltinStablehloMaximum:
+      return ynn_binary_max;
+    case kTfLiteBuiltinMinimum:
+    case kTfLiteBuiltinStablehloMinimum:
+      return ynn_binary_min;
+    case kTfLiteBuiltinMul:
+    case kTfLiteBuiltinStablehloMultiply:
+      return ynn_binary_multiply;
+    case kTfLiteBuiltinPow:
+    case kTfLiteBuiltinStablehloPower:
+      return ynn_binary_pow;
+    case kTfLiteBuiltinSub:
+    case kTfLiteBuiltinStablehloSubtract:
+      return ynn_binary_subtract;
+    case kTfLiteBuiltinSquaredDifference:
+      return ynn_binary_squared_difference;
+    case kTfLiteBuiltinPrelu:
+      return ynn_binary_leaky_relu;
+    default:
+      return ynn_binary_invalid;
+  }
+}
+
+ynn_reduce_operator GetYnnReduceOperator(int builtin_code) {
+  switch (builtin_code) {
+    case kTfLiteBuiltinSum:
+    case kTfLiteBuiltinMean:
+      return ynn_reduce_sum;
+    case kTfLiteBuiltinReduceMin:
+      return ynn_reduce_min;
+    case kTfLiteBuiltinReduceMax:
+      return ynn_reduce_max;
+    default:
+      return ynn_reduce_invalid;
+  }
+}
+
+bool IsUnaryOp(int builtin_code) {
+  return GetYnnUnaryOperator(builtin_code) != ynn_unary_invalid ||
+         builtin_code == kTfLiteBuiltinGelu ||
+         builtin_code == kTfLiteBuiltinElu ||
+         builtin_code == kTfLiteBuiltinLeakyRelu ||
+         builtin_code == kTfLiteBuiltinHardSwish ||
+         builtin_code == kTfLiteBuiltinRelu ||
+         builtin_code == kTfLiteBuiltinRelu6 ||
+         builtin_code == kTfLiteBuiltinReluN1To1 ||
+         builtin_code == kTfLiteBuiltinRelu0To1;
+}
+
+bool IsBinaryOp(int builtin_code) {
+  return GetYnnBinaryOperator(builtin_code) != ynn_binary_invalid;
+}
+
+bool IsStablehloOp(int builtin_code) {
+  switch (builtin_code) {
+    case kTfLiteBuiltinStablehloLogistic:
+    case kTfLiteBuiltinStablehloAdd:
+    case kTfLiteBuiltinStablehloDivide:
+    case kTfLiteBuiltinStablehloMultiply:
+    case kTfLiteBuiltinStablehloMaximum:
+    case kTfLiteBuiltinStablehloAbs:
+    case kTfLiteBuiltinStablehloCosine:
+    case kTfLiteBuiltinStablehloExponential:
+    case kTfLiteBuiltinStablehloFloor:
+    case kTfLiteBuiltinStablehloLog:
+    case kTfLiteBuiltinStablehloMinimum:
+    case kTfLiteBuiltinStablehloNegate:
+    case kTfLiteBuiltinStablehloPower:
+    case kTfLiteBuiltinStablehloRsqrt:
+    case kTfLiteBuiltinStablehloSubtract:
+    case kTfLiteBuiltinStablehloTanh:
+    case kTfLiteBuiltinStablehloConvert:
+    case kTfLiteBuiltinStablehloClamp:
+      return true;
+    default:
+      return false;
+  }
+}
+
+bool IsQuantized(const TfLiteTensor& tensor) {
+  if (tensor.type == kTfLiteFloat32 || tensor.type == kTfLiteFloat16 ||
+      tensor.type == kTfLiteBFloat16 || tensor.type == kTfLiteFloat64 ||
+      tensor.type == kTfLiteFloat8E4M3FN || tensor.type == kTfLiteFloat8E5M2) {
+    return false;
+  }
+  return tensor.quantization.type != kTfLiteNoQuantization ||
+         tensor.params.scale != 0.0f;
+}
+
+bool IsSupportedQuantization(const TfLiteTensor& tensor,
+                             bool allow_per_channel) {
+  if (!IsQuantized(tensor)) return true;
+  if (tensor.quantization.type != kTfLiteAffineQuantization) return false;
+  const auto* params =
+      static_cast<const TfLiteAffineQuantization*>(tensor.quantization.params);
+  if (!params) return false;
+  if (!params->scale) return false;
+  if (params->scale->size != 1 && !allow_per_channel) return false;
+  if (tensor.type == kTfLiteUInt8) {
+    if (!params->zero_point) return false;
+    if (params->zero_point->size != params->scale->size &&
+        params->zero_point->size != 1) {
+      return false;
+    }
+  }
+  return true;
+}
+
+size_t YnnTypeElementCount(ynn_type type) {
+  switch (type) {
+    case ynn_type_int2:
+    case ynn_type_uint2:
+      return 4;
+    case ynn_type_int4:
+    case ynn_type_uint4:
+      return 2;
+    default:
+      return 1;
+  }
+}
+
+bool IsTensorSupported(const TfLiteTensor& tensor, bool allow_per_channel) {
+  ynn_type type = GetYnnType(tensor.type);
+  if (type == ynn_type_invalid) return false;
+  size_t element_count = YnnTypeElementCount(type);
+  if (element_count > 1) {
+    if (tensor.dims == nullptr) return false;
+    size_t dense_dim =
+        tensor.dims->size == 0 ? 1 : tensor.dims->data[tensor.dims->size - 1];
+    if (dense_dim % element_count != 0) return false;
+  }
+  return IsSupportedQuantization(tensor, allow_per_channel);
+}
+
+bool QuantizationParamsEqual(const TfLiteTensor& tensor1,
+                             const TfLiteTensor& tensor2) {
+  if (IsQuantized(tensor1) != IsQuantized(tensor2)) return false;
+  if (!IsQuantized(tensor1)) return true;
+  return tensor1.params.scale == tensor2.params.scale &&
+         tensor1.params.zero_point == tensor2.params.zero_point;
+}
+
+bool IsActivationSupported(TfLiteFusedActivation activation,
+                           TfLiteType output_type) {
+  if (activation == kTfLiteActNone) return true;
+
+  return (activation == kTfLiteActRelu || activation == kTfLiteActRelu6 ||
+          activation == kTfLiteActReluN1To1 || activation == kTfLiteActTanh ||
+          activation == kTfLiteActSigmoid);
+}
+
+TfLiteFusedActivation GetFusedActivation(const TfLiteRegistration* registration,
+                                         const TfLiteNode* node) {
+  TfLiteFusedActivation activation = kTfLiteActNone;
+  switch (registration->builtin_code) {
+    case kTfLiteBuiltinAdd: {
+      const auto* params =
+          reinterpret_cast<const TfLiteAddParams*>(node->builtin_data);
+      if (params) activation = params->activation;
+      break;
+    }
+    case kTfLiteBuiltinSub: {
+      const auto* params =
+          reinterpret_cast<const TfLiteSubParams*>(node->builtin_data);
+      if (params) activation = params->activation;
+      break;
+    }
+    case kTfLiteBuiltinMul: {
+      const auto* params =
+          reinterpret_cast<const TfLiteMulParams*>(node->builtin_data);
+      if (params) activation = params->activation;
+      break;
+    }
+    case kTfLiteBuiltinDiv: {
+      const auto* params =
+          reinterpret_cast<const TfLiteDivParams*>(node->builtin_data);
+      if (params) activation = params->activation;
+      break;
+    }
+    case kTfLiteBuiltinMaxPool2d:
+    case kTfLiteBuiltinAveragePool2d: {
+      const auto* params =
+          reinterpret_cast<const TfLitePoolParams*>(node->builtin_data);
+      if (params) activation = params->activation;
+      break;
+    }
+    default:
+      break;
+  }
+  return activation;
+}
+
+TfLiteStatus DefineScalarConstant(TfLiteContext* context,
+                                  ynn_subgraph_t subgraph, ynn_type type,
+                                  double value, uint32_t* id_out) {
+  TF_LITE_ENSURE(context, type != ynn_type_invalid);
+  switch (type) {
+    case ynn_type_fp32: {
+      float f_val = static_cast<float>(value);
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, type, 0, nullptr, &f_val,
+                                          YNN_VALUE_FLAG_COPY_DATA, id_out),
+                        ynn_status_success);
+      return kTfLiteOk;
+    }
+    case ynn_type_int32: {
+      int32_t i_val = static_cast<int32_t>(value);
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, type, 0, nullptr, &i_val,
+                                          YNN_VALUE_FLAG_COPY_DATA, id_out),
+                        ynn_status_success);
+      return kTfLiteOk;
+    }
+    case ynn_type_int8: {
+      int8_t i8_val = static_cast<int8_t>(value);
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, type, 0, nullptr, &i8_val,
+                                          YNN_VALUE_FLAG_COPY_DATA, id_out),
+                        ynn_status_success);
+      return kTfLiteOk;
+    }
+    case ynn_type_uint8: {
+      uint8_t u8_val = static_cast<uint8_t>(value);
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, type, 0, nullptr, &u8_val,
+                                          YNN_VALUE_FLAG_COPY_DATA, id_out),
+                        ynn_status_success);
+      return kTfLiteOk;
+    }
+    case ynn_type_fp64: {
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, type, 0, nullptr, &value,
+                                          YNN_VALUE_FLAG_COPY_DATA, id_out),
+                        ynn_status_success);
+      return kTfLiteOk;
+    }
+    case ynn_type_fp16: {
+      tflite::half f16_val(static_cast<float>(value));
+      uint16_t bits = f16_val.to_bits();
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, type, 0, nullptr, &bits,
+                                          YNN_VALUE_FLAG_COPY_DATA, id_out),
+                        ynn_status_success);
+      return kTfLiteOk;
+    }
+    case ynn_type_bf16: {
+      uint16_t bits = FloatToBfloat16(static_cast<float>(value));
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, type, 0, nullptr, &bits,
+                                          YNN_VALUE_FLAG_COPY_DATA, id_out),
+                        ynn_status_success);
+      return kTfLiteOk;
+    }
+    default:
+      TF_LITE_ENSURE(context, false);
+  }
+}
+
+TfLiteStatus DefineQuantizationParams(TfLiteContext* context,
+                                      ynn_subgraph_t subgraph,
+                                      const TfLiteTensor& tensor,
+                                      uint32_t* scale_id, uint32_t* zp_id,
+                                      int32_t zp_offset) {
+  if (!IsQuantized(tensor)) {
+    *scale_id = YNN_INVALID_VALUE_ID;
+    *zp_id = YNN_INVALID_VALUE_ID;
+    return kTfLiteOk;
+  }
+  const auto* params =
+      static_cast<const TfLiteAffineQuantization*>(tensor.quantization.params);
+  if (params && params->scale && params->scale->size > 1) {
+    TF_LITE_ENSURE(context, params->zero_point != nullptr);
+    TF_LITE_ENSURE(context, params->zero_point->size == params->scale->size ||
+                                params->zero_point->size == 1);
+
+    size_t size = params->scale->size;
+    size_t dims[] = {size};
+
+    TF_LITE_ENSURE_EQ(
+        context,
+        ynn_define_tensor(subgraph, ynn_type_fp32, 1, dims, params->scale->data,
+                          YNN_VALUE_FLAG_COPY_DATA, scale_id),
+        ynn_status_success);
+
+    std::vector<int32_t> adjusted_zp(size);
+    for (size_t i = 0; i < size; ++i) {
+      int32_t zp = params->zero_point->size == 1 ? params->zero_point->data[0]
+                                                 : params->zero_point->data[i];
+      adjusted_zp[i] = zp + zp_offset;
+    }
+    bool all_zero = std::all_of(adjusted_zp.begin(), adjusted_zp.end(),
+                                [](int32_t zp) { return zp == 0; });
+    if (all_zero) {
+      *zp_id = YNN_INVALID_VALUE_ID;
+    } else {
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, ynn_type_int32, 1, dims,
+                                          adjusted_zp.data(),
+                                          YNN_VALUE_FLAG_COPY_DATA, zp_id),
+                        ynn_status_success);
+    }
+  } else {
+    float scale = 0.0f;
+    int32_t zp = 0;
+    if (params && params->scale) {
+      scale = params->scale->data[0];
+      zp = params->zero_point ? params->zero_point->data[0] : 0;
+    } else {
+      TFLITE_LOG_PROD(
+          TFLITE_LOG_WARNING,
+          "DefineQuantizationParams: missing or invalid affine quantization "
+          "params for tensor, using legacy params. scale: %f, zp: %d",
+          tensor.params.scale, tensor.params.zero_point);
+      scale = tensor.params.scale;
+      zp = tensor.params.zero_point;
+      TF_LITE_ENSURE_MSG(context, scale != 0.0f,
+                         "DefineQuantizationParams: scale is 0.0");
+    }
+
+    zp += zp_offset;
+
+    TF_LITE_ENSURE_EQ(
+        context,
+        ynn_define_tensor(subgraph, ynn_type_fp32, 0, nullptr, &scale,
+                          YNN_VALUE_FLAG_COPY_DATA, scale_id),
+        ynn_status_success);
+
+    if (zp == 0) {
+      *zp_id = YNN_INVALID_VALUE_ID;
+    } else {
+      TF_LITE_ENSURE_EQ(context,
+                        ynn_define_tensor(subgraph, ynn_type_int32, 0, nullptr,
+                                          &zp, YNN_VALUE_FLAG_COPY_DATA, zp_id),
+                        ynn_status_success);
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus ApplyClamp(TfLiteContext* context, ynn_subgraph_t subgraph,
+                        double min_val, double max_val, uint32_t input_id,
+                        uint32_t& output_id, int output_tensor_index,
+                        ynn_type internal_type) {
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+  const bool is_quantized =
+      IsQuantized(output_tensor) &&
+      (internal_type == ynn_type_int8 || internal_type == ynn_type_uint8);
+
+  auto quantize_val = [&](double val) -> double {
+    if (!is_quantized) return val;
+    const float scale = output_tensor.params.scale;
+    const int32_t zero_point = output_tensor.params.zero_point;
+    if (internal_type == ynn_type_int8) {
+      return QuantizeValue<int8_t>(val, scale, zero_point);
+    } else {
+      return QuantizeValue<uint8_t>(val, scale, zero_point);
+    }
+  };
+
+  uint32_t current_input_id = input_id;
+
+  uint32_t min_const_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DefineScalarConstant(
+      context, subgraph, internal_type, quantize_val(min_val), &min_const_id));
+
+  uint32_t max_output_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_EQ(
+      context,
+      ynn_define_binary(subgraph, ynn_binary_max, current_input_id,
+                        min_const_id, &max_output_id, 0),
+      ynn_status_success);
+
+  current_input_id = max_output_id;
+
+  uint32_t max_const_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DefineScalarConstant(
+      context, subgraph, internal_type, quantize_val(max_val), &max_const_id));
+
+  TF_LITE_ENSURE_EQ(
+      context,
+      ynn_define_binary(subgraph, ynn_binary_min, current_input_id,
+                        max_const_id, &output_id, 0),
+      ynn_status_success);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus ApplyActivation(TfLiteContext* context, ynn_subgraph_t subgraph,
+                             TfLiteFusedActivation activation,
+                             uint32_t input_id, uint32_t& output_id,
+                             int output_tensor_index, ynn_type internal_type) {
+  if (activation == kTfLiteActNone) {
+    return kTfLiteOk;
+  }
+
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+  bool is_quantized =
+      IsQuantized(output_tensor) &&
+      (internal_type == ynn_type_int8 || internal_type == ynn_type_uint8);
+
+  auto quantize_val = [&](double val) -> double {
+    if (!is_quantized) return val;
+    float scale = output_tensor.params.scale;
+    int32_t zero_point = output_tensor.params.zero_point;
+    if (internal_type == ynn_type_int8) {
+      return QuantizeValue<int8_t>(val, scale, zero_point);
+    } else {
+      return QuantizeValue<uint8_t>(val, scale, zero_point);
+    }
+  };
+
+  uint32_t current_input_id = input_id;
+
+  if (activation == kTfLiteActRelu) {
+    uint32_t min_const_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_STATUS(DefineScalarConstant(
+        context, subgraph, internal_type, quantize_val(0.0), &min_const_id));
+    TF_LITE_ENSURE_EQ(context,
+                      ynn_define_binary(subgraph, ynn_binary_max, input_id,
+                                        min_const_id, &output_id, 0),
+                      ynn_status_success);
+  } else if (activation == kTfLiteActRelu6 ||
+             activation == kTfLiteActReluN1To1) {
+    double min_val = (activation == kTfLiteActReluN1To1) ? -1.0 : 0.0;
+    double max_val = (activation == kTfLiteActRelu6) ? 6.0 : 1.0;
+    return ApplyClamp(context, subgraph, min_val, max_val, input_id, output_id,
+                      output_tensor_index, internal_type);
+  } else if (activation == kTfLiteActTanh) {
+    TF_LITE_ENSURE_EQ(context,
+                      ynn_define_unary(subgraph, ynn_unary_tanh,
+                                       current_input_id, &output_id, 0),
+                      ynn_status_success);
+  } else if (activation == kTfLiteActSigmoid) {
+    TF_LITE_ENSURE_EQ(context,
+                      ynn_define_unary(subgraph, ynn_unary_sigmoid,
+                                       current_input_id, &output_id, 0),
+                      ynn_status_success);
+  } else {
+    TF_LITE_ENSURE(context, false);
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DequantizeIfNeeded(TfLiteContext* context, ynn_subgraph_t subgraph,
+                                TensorToValueIdMap& tensor_to_value_id,
+                                int tensor_index, uint32_t val_id,
+                                uint32_t* float_val_id) {
+  const TfLiteTensor& tensor = context->tensors[tensor_index];
+  if (IsQuantized(tensor)) {
+    uint32_t scale_id = YNN_INVALID_VALUE_ID;
+    uint32_t zp_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_STATUS(
+        DefineQuantizationParams(context, subgraph, tensor, &scale_id, &zp_id));
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_dequantize(
+        subgraph, val_id, zp_id, scale_id, ynn_type_fp32, float_val_id, 0));
+  } else {
+    *float_val_id = val_id;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus Quantize(TfLiteContext* context, ynn_subgraph_t subgraph,
+                      TensorToValueIdMap& tensor_to_value_id, int tensor_index,
+                      uint32_t float_val_id, uint32_t quant_val_id) {
+  const TfLiteTensor& tensor = context->tensors[tensor_index];
+  uint32_t scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t zp_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(
+      DefineQuantizationParams(context, subgraph, tensor, &scale_id, &zp_id));
+  ynn_type ynn_type = GetYnnType(tensor.type);
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_quantize(
+      subgraph, float_val_id, ynn_type, zp_id, scale_id, &quant_val_id, 0));
+  return kTfLiteOk;
+}
+
+uint32_t GetOrCreateValueId(TfLiteContext* context, ynn_subgraph_t subgraph,
+                            TensorToValueIdMap& mapping, int tensor_index) {
+  if (tensor_index < 0 || tensor_index >= context->tensors_size) {
+    return YNN_INVALID_VALUE_ID;
+  }
+
+  auto it = mapping.find(tensor_index);
+  if (it != mapping.end()) return it->second;
+
+  const TfLiteTensor& tensor = context->tensors[tensor_index];
+  uint32_t val_id = YNN_INVALID_VALUE_ID;
+  ynn_type ynn_type = GetYnnType(tensor.type);
+
+  if (ynn_type == ynn_type_invalid) {
+    return YNN_INVALID_VALUE_ID;
+  }
+  if (tensor.dims->size > YNN_MAX_TENSOR_RANK) {
+    return YNN_INVALID_VALUE_ID;
+  }
+
+  const void* data = nullptr;
+  uint32_t flags = 0;
+  size_t dims_data[YNN_MAX_TENSOR_RANK];
+  const size_t* dims = nullptr;
+
+  if (tensor.allocation_type == kTfLiteMmapRo) {
+    data = tensor.data.raw;
+    std::copy_n(tensor.dims->data, tensor.dims->size, dims_data);
+    dims = dims_data;
+  }
+
+  ynn_status status = ynn_define_tensor(subgraph, ynn_type, tensor.dims->size,
+                                        dims, data, flags, &val_id);
+  if (status != ynn_status_success) return YNN_INVALID_VALUE_ID;
+  mapping[tensor_index] = val_id;
+  return val_id;
+}
+
+TfLiteStatus ImplementMutualBroadcasting(TfLiteContext* context,
+                                         ynn_subgraph_t subgraph, int rank_a,
+                                         int rank_b, int exclude_a,
+                                         int exclude_b, uint32_t& current_a_id,
+                                         uint32_t& current_b_id) {
+  int max_rank = std::max(rank_a, rank_b);
+
+  uint32_t input_a_id = current_a_id;
+  uint32_t input_b_id = current_b_id;
+
+  // 1. Expand A if needed.
+  uint32_t expanded_a_id = input_a_id;
+  if (max_rank > rank_a) {
+    int diff = max_rank - rank_a;
+    int32_t axes[YNN_MAX_TENSOR_RANK];
+    std::iota(axes, axes + diff, 0);
+    expanded_a_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_expand_dims(
+        subgraph, diff, axes, input_a_id, &expanded_a_id, 0));
+  }
+
+  // 2. Expand B if needed.
+  uint32_t expanded_b_id = input_b_id;
+  if (max_rank > rank_b) {
+    int diff = max_rank - rank_b;
+    int32_t axes[YNN_MAX_TENSOR_RANK];
+    std::iota(axes, axes + diff, 0);
+    expanded_b_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_expand_dims(
+        subgraph, diff, axes, input_b_id, &expanded_b_id, 0));
+  }
+
+  uint32_t broadcasted_a_id = expanded_a_id;
+  if (max_rank > exclude_a) {
+    broadcasted_a_id = YNN_INVALID_VALUE_ID;
+    int num_axes = max_rank - exclude_a;
+    int32_t axes[YNN_MAX_TENSOR_RANK];
+    std::iota(axes, axes + num_axes, 0);
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_broadcast_like(subgraph, num_axes, axes, expanded_a_id,
+                                  expanded_b_id, &broadcasted_a_id, 0));
+  }
+
+  uint32_t broadcasted_b_id = expanded_b_id;
+  if (max_rank > exclude_b) {
+    broadcasted_b_id = YNN_INVALID_VALUE_ID;
+    int num_axes = max_rank - exclude_b;
+    int32_t axes[YNN_MAX_TENSOR_RANK];
+    std::iota(axes, axes + num_axes, 0);
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_broadcast_like(subgraph, num_axes, axes, expanded_b_id,
+                                  expanded_a_id, &broadcasted_b_id, 0));
+  }
+
+  current_a_id = broadcasted_a_id;
+  current_b_id = broadcasted_b_id;
+  return kTfLiteOk;
+}
+
+TfLiteStatus DefineYnnStencil(TfLiteContext* context, ynn_subgraph_t subgraph,
+                              const TfLiteTensor& input_tensor,
+                              uint32_t input_id, size_t filter_height,
+                              size_t filter_width, size_t stride_height,
+                              size_t stride_width, size_t dilation_height,
+                              size_t dilation_width, TfLitePadding padding,
+                              float padding_value, uint32_t* stencil_id) {
+  uint32_t padding_id = YNN_INVALID_VALUE_ID;
+  if (padding == kTfLitePaddingSame) {
+    ynn_type padding_type = GetYnnType(input_tensor.type);
+    TF_LITE_ENSURE(context, padding_type != ynn_type_invalid);
+    TF_LITE_ENSURE_STATUS(DefineScalarConstant(context, subgraph, padding_type,
+                                               padding_value, &padding_id));
+  }
+
+  const int32_t stencil_axes[] = {1, 2};
+  const int32_t new_axes[] = {-3, -2};
+  const size_t stencil_dims[] = {filter_height, filter_width};
+  const size_t stencil_strides[] = {stride_height, stride_width};
+  const size_t stencil_dilations[] = {dilation_height, dilation_width};
+
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_stencil_copy(
+      subgraph, /*num_stencils=*/2, stencil_axes, new_axes, stencil_dims,
+      stencil_strides, stencil_dilations, input_id, padding_id, stencil_id,
+      /*flags=*/0));
+
+  return kTfLiteOk;
+}
+
+}  // namespace ynnpack
+}  // namespace tflite

tensorflow/lite/delegates/ynnpack/utils.h

@@ -0,0 +1,182 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_UTILS_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_UTILS_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+
+#define TF_LITE_ENSURE_YNN_STATUS(x) \
+  TF_LITE_ENSURE(context, (x) == ynn_status_success)
+
+namespace tflite {
+namespace ynnpack {
+
+using TensorToValueIdMap = absl::flat_hash_map<int, uint32_t>;
+
+struct NodeInfo {
+  int node_index;
+  int builtin_code;
+  std::vector<int> inputs;
+  std::vector<int> outputs;
+  TfLiteFusedActivation activation;
+};
+
+// Generic helpers
+ynn_type GetYnnType(TfLiteType type);
+size_t YnnTypeElementCount(ynn_type type);
+ynn_unary_operator GetYnnUnaryOperator(int builtin_code);
+ynn_binary_operator GetYnnBinaryOperator(int builtin_code);
+ynn_reduce_operator GetYnnReduceOperator(int builtin_code);
+bool IsUnaryOp(int builtin_code);
+bool IsBinaryOp(int builtin_code);
+bool IsStablehloOp(int builtin_code);
+bool IsQuantized(const TfLiteTensor& tensor);
+bool IsSupportedQuantization(const TfLiteTensor& tensor,
+                             bool allow_per_channel = false);
+bool IsTensorSupported(const TfLiteTensor& tensor,
+                       bool allow_per_channel = false);
+bool QuantizationParamsEqual(const TfLiteTensor& tensor1,
+                             const TfLiteTensor& tensor2);
+bool IsActivationSupported(TfLiteFusedActivation activation,
+                           TfLiteType output_type);
+TfLiteFusedActivation GetFusedActivation(const TfLiteRegistration* registration,
+                                         const TfLiteNode* node);
+
+TfLiteStatus GetTfLiteTensorValueAsDouble(TfLiteContext* context,
+                                          const TfLiteTensor& tensor, int index,
+                                          double* value);
+
+// YNNPACK graph building helpers
+TfLiteStatus DefineYnnStencil(TfLiteContext* context, ynn_subgraph_t subgraph,
+                              const TfLiteTensor& input_tensor,
+                              uint32_t input_id, size_t filter_height,
+                              size_t filter_width, size_t stride_height,
+                              size_t stride_width, size_t dilation_height,
+                              size_t dilation_width, TfLitePadding padding,
+                              float padding_value, uint32_t* stencil_id);
+
+TfLiteStatus DefineScalarConstant(TfLiteContext* context,
+                                  ynn_subgraph_t subgraph, ynn_type type,
+                                  double value, uint32_t* id_out);
+
+TfLiteStatus DefineQuantizationParams(TfLiteContext* context,
+                                      ynn_subgraph_t subgraph,
+                                      const TfLiteTensor& tensor,
+                                      uint32_t* scale_id, uint32_t* zp_id,
+                                      int32_t zp_offset = 0);
+
+TfLiteStatus ApplyActivation(TfLiteContext* context, ynn_subgraph_t subgraph,
+                             TfLiteFusedActivation activation,
+                             uint32_t input_id, uint32_t& output_id,
+                             int output_tensor_index, ynn_type internal_type);
+
+TfLiteStatus ApplyClamp(TfLiteContext* context, ynn_subgraph_t subgraph,
+                        double min_val, double max_val, uint32_t input_id,
+                        uint32_t& output_id, int output_tensor_index,
+                        ynn_type internal_type);
+
+TfLiteStatus DequantizeIfNeeded(TfLiteContext* context, ynn_subgraph_t subgraph,
+                                TensorToValueIdMap& tensor_to_value_id,
+                                int tensor_index, uint32_t val_id,
+                                uint32_t* float_val_id);
+
+TfLiteStatus Quantize(TfLiteContext* context, ynn_subgraph_t subgraph,
+                      TensorToValueIdMap& tensor_to_value_id, int tensor_index,
+                      uint32_t float_val_id, uint32_t quant_val_id);
+
+uint32_t GetOrCreateValueId(TfLiteContext* context, ynn_subgraph_t subgraph,
+                            TensorToValueIdMap& mapping, int tensor_index);
+
+TfLiteStatus ImplementMutualBroadcasting(TfLiteContext* context,
+                                         ynn_subgraph_t subgraph, int rank_a,
+                                         int rank_b, int exclude_a,
+                                         int exclude_b, uint32_t& current_a_id,
+                                         uint32_t& current_b_id);
+
+// Template helper for decomposed unary nodes
+template <typename F>
+TfLiteStatus DefineDecomposedUnaryNode(TfLiteContext* context,
+                                       ynn_subgraph_t subgraph,
+                                       TensorToValueIdMap& tensor_to_value_id,
+                                       const NodeInfo& node,
+                                       F&& define_float_subgraph) {
+  TF_LITE_ENSURE_EQ(context, node.inputs.size(), 1);
+  TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1);
+
+  int input_tensor_index = node.inputs[0];
+  int output_tensor_index = node.outputs[0];
+
+  uint32_t input_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, input_tensor_index);
+  TF_LITE_ENSURE(context, input_val_id != YNN_INVALID_VALUE_ID);
+  uint32_t output_val_id = GetOrCreateValueId(
+      context, subgraph, tensor_to_value_id, output_tensor_index);
+  TF_LITE_ENSURE(context, output_val_id != YNN_INVALID_VALUE_ID);
+
+  const TfLiteTensor& output_tensor = context->tensors[output_tensor_index];
+  bool is_output_quantized = IsQuantized(output_tensor);
+  TfLiteFusedActivation activation = node.activation;
+
+  uint32_t float_input_val_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DequantizeIfNeeded(
+      context, subgraph, tensor_to_value_id, input_tensor_index, input_val_id,
+      &float_input_val_id));
+
+  ynn_type internal_type = IsQuantized(output_tensor)
+                               ? ynn_type_fp32
+                               : GetYnnType(output_tensor.type);
+  uint32_t float_output_val_id = YNN_INVALID_VALUE_ID;
+
+  if (!is_output_quantized && activation == kTfLiteActNone) {
+    float_output_val_id = output_val_id;
+  }
+
+  TF_LITE_ENSURE_STATUS(
+      define_float_subgraph(float_input_val_id, float_output_val_id));
+
+  uint32_t active_output_val_id = float_output_val_id;
+
+  if (activation != kTfLiteActNone) {
+    uint32_t activation_output_val_id = output_val_id;
+    if (is_output_quantized) {
+      activation_output_val_id = YNN_INVALID_VALUE_ID;
+    }
+    TF_LITE_ENSURE_STATUS(ApplyActivation(
+        context, subgraph, activation, float_output_val_id,
+        activation_output_val_id, output_tensor_index, internal_type));
+    active_output_val_id = activation_output_val_id;
+  }
+
+  if (is_output_quantized) {
+    TF_LITE_ENSURE_STATUS(Quantize(context, subgraph, tensor_to_value_id,
+                                   output_tensor_index, active_output_val_id,
+                                   output_val_id));
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_UTILS_H_

tensorflow/lite/delegates/ynnpack/ynnpack_delegate.cc

@@ -0,0 +1,498 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/ynnpack/ynnpack_delegate.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "ynnpack/include/ynnpack.h"  // from @XNNPACK
+#include "slinky/base/thread_pool.h"  // from @slinky
+#include "slinky/base/thread_pool_impl.h"  // from @slinky
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/utils/simple_delegate.h"
+#include "tensorflow/lite/delegates/ynnpack/copy.h"
+#include "tensorflow/lite/delegates/ynnpack/dot.h"
+#include "tensorflow/lite/delegates/ynnpack/elementwise.h"
+#include "tensorflow/lite/delegates/ynnpack/pooling.h"
+#include "tensorflow/lite/delegates/ynnpack/reduction.h"
+#include "tensorflow/lite/delegates/ynnpack/softmax.h"
+#include "tensorflow/lite/delegates/ynnpack/utils.h"
+
+namespace tflite {
+namespace ynnpack {
+
+class YNNPackDelegateKernel : public SimpleDelegateKernelInterface {
+ public:
+  explicit YNNPackDelegateKernel(const TfLiteYNNPackDelegateOptions& options,
+                                 slinky::thread_pool* thread_pool)
+      : options_(options),
+        thread_pool_(thread_pool),
+        subgraph_(nullptr),
+        runtime_(nullptr) {}
+
+  ~YNNPackDelegateKernel() override {
+    if (runtime_) ynn_delete_runtime(runtime_);
+    if (subgraph_) ynn_delete_subgraph(subgraph_);
+  }
+
+  TfLiteStatus BuildSubgraphAndRuntime(TfLiteContext* context) {
+    if (runtime_) {
+      ynn_delete_runtime(runtime_);
+      runtime_ = nullptr;
+    }
+    if (subgraph_) {
+      ynn_delete_subgraph(subgraph_);
+      subgraph_ = nullptr;
+    }
+    tensor_to_value_id_.clear();
+    inputs_.clear();
+
+    outputs_.clear();
+    input_shapes_.clear();
+
+    int external_value_ids =
+        input_tensor_indices_.size() + output_tensor_indices_.size();
+    uint32_t subgraph_flags = 0;
+    if (options_.fast_math) {
+      subgraph_flags |= YNN_FLAG_FAST_MATH;
+    }
+    if (options_.consistent_arithmetic) {
+      subgraph_flags |= YNN_FLAG_CONSISTENT_ARITHMETIC;
+    }
+    if (options_.no_excess_precision) {
+      subgraph_flags |= YNN_FLAG_NO_EXCESS_PRECISION;
+    }
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_create_subgraph(external_value_ids, subgraph_flags, &subgraph_));
+
+    int next_external_id = 0;
+
+    // Define input tensors of the partition as external inputs.
+    input_shapes_.resize(input_tensor_indices_.size());
+    for (size_t i = 0; i < input_tensor_indices_.size(); ++i) {
+      int tensor_index = input_tensor_indices_[i];
+      const TfLiteTensor& tensor = context->tensors[tensor_index];
+      uint32_t val_id = next_external_id++;
+
+      ynn_type ynn_type = GetYnnType(tensor.type);
+      TF_LITE_ENSURE_MSG(context, ynn_type != ynn_type_invalid,
+                         "Unsupported type %d for input tensor %d in Build",
+                         tensor.type, tensor_index);
+
+      TF_LITE_ENSURE_MSG(context, tensor.dims->size <= YNN_MAX_TENSOR_RANK,
+                         "Tensor %d rank %d exceeds max %d", tensor_index,
+                         tensor.dims->size, YNN_MAX_TENSOR_RANK);
+      if (tensor.allocation_type == kTfLiteMmapRo) {
+        size_t dims[YNN_MAX_TENSOR_RANK];
+        std::copy_n(tensor.dims->data, tensor.dims->size, dims);
+        TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(
+            subgraph_, ynn_type, tensor.dims->size, dims, tensor.data.raw,
+            /*flags=*/0, &val_id));
+      } else {
+        size_t dims[YNN_MAX_TENSOR_RANK];
+        const size_t* dims_ptr = nullptr;
+        if (options_.static_shape) {
+          std::copy_n(tensor.dims->data, tensor.dims->size, dims);
+          dims_ptr = dims;
+        }
+        TF_LITE_ENSURE_YNN_STATUS(
+            ynn_define_tensor(subgraph_, ynn_type, tensor.dims->size, dims_ptr,
+                              nullptr, YNN_VALUE_FLAG_EXTERNAL_INPUT, &val_id));
+        inputs_.push_back({tensor_index, val_id});
+      }
+      input_shapes_[i].assign(tensor.dims->data,
+                              tensor.dims->data + tensor.dims->size);
+
+      tensor_to_value_id_[tensor_index] = val_id;
+    }
+
+    // Define output tensors of the partition as external outputs.
+    outputs_.reserve(output_tensor_indices_.size());
+    for (size_t i = 0; i < output_tensor_indices_.size(); ++i) {
+      int tensor_index = output_tensor_indices_[i];
+      const TfLiteTensor& tensor = context->tensors[tensor_index];
+      uint32_t val_id = next_external_id++;
+
+      ynn_type ynn_type = GetYnnType(tensor.type);
+      TF_LITE_ENSURE_MSG(context, ynn_type != ynn_type_invalid,
+                         "Unsupported type %d for output tensor %d in Build",
+                         tensor.type, tensor_index);
+
+      TF_LITE_ENSURE_MSG(context, tensor.dims->size <= YNN_MAX_TENSOR_RANK,
+                         "Tensor %d rank %d exceeds max %d", tensor_index,
+                         tensor.dims->size, YNN_MAX_TENSOR_RANK);
+      size_t dims[YNN_MAX_TENSOR_RANK];
+      std::copy_n(tensor.dims->data, tensor.dims->size, dims);
+
+      TF_LITE_ENSURE_YNN_STATUS(
+          ynn_define_tensor(subgraph_, ynn_type, tensor.dims->size, dims,
+                            nullptr, YNN_VALUE_FLAG_EXTERNAL_OUTPUT, &val_id));
+      outputs_.push_back({tensor_index, val_id});
+
+      tensor_to_value_id_[tensor_index] = val_id;
+    }
+
+    // Now define internal nodes.
+    for (const auto& node : nodes_info_) {
+      if (IsUnaryOp(node.builtin_code)) {
+        TF_LITE_ENSURE_STATUS(
+            DefineUnaryNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinSoftmax) {
+        TF_LITE_ENSURE_STATUS(
+            DefineSoftmaxNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinLogSoftmax) {
+        TF_LITE_ENSURE_STATUS(DefineLogSoftmaxNode(context, subgraph_,
+                                                   tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinMaxPool2d) {
+        TF_LITE_ENSURE_STATUS(
+            DefineMaxPool2DNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinAveragePool2d) {
+        TF_LITE_ENSURE_STATUS(DefineAveragePool2DNode(
+            context, subgraph_, tensor_to_value_id_, node));
+      } else if (IsBinaryOp(node.builtin_code)) {
+        TF_LITE_ENSURE_STATUS(
+            DefineBinaryNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinTranspose) {
+        TF_LITE_ENSURE_STATUS(
+            DefineTransposeNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinGather) {
+        TF_LITE_ENSURE_STATUS(
+            DefineGatherNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinGatherNd) {
+        TF_LITE_ENSURE_STATUS(
+            DefineGatherNdNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinSlice ||
+                 node.builtin_code == kTfLiteBuiltinStridedSlice) {
+        TF_LITE_ENSURE_STATUS(
+            DefineSliceNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinConcatenation) {
+        TF_LITE_ENSURE_STATUS(DefineConcatenationNode(
+            context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinReshape) {
+        TF_LITE_ENSURE_STATUS(
+            DefineReshapeNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinExpandDims) {
+        TF_LITE_ENSURE_STATUS(DefineExpandDimsNode(context, subgraph_,
+                                                   tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinPad ||
+                 node.builtin_code == kTfLiteBuiltinPadv2) {
+        TF_LITE_ENSURE_STATUS(
+            DefinePadNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinSplit) {
+        TF_LITE_ENSURE_STATUS(
+            DefineSplitNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinSpaceToDepth) {
+        TF_LITE_ENSURE_STATUS(DefineSpaceToDepthNode(
+            context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinDepthToSpace) {
+        TF_LITE_ENSURE_STATUS(DefineDepthToSpaceNode(
+            context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinBatchMatmul) {
+        TF_LITE_ENSURE_STATUS(DefineBatchMatMulNode(context, subgraph_,
+                                                    tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinFullyConnected) {
+        TF_LITE_ENSURE_STATUS(DefineFullyConnectedNode(
+            context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinConv2d) {
+        TF_LITE_ENSURE_STATUS(
+            DefineConvNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinDepthwiseConv2d) {
+        TF_LITE_ENSURE_STATUS(DefineDepthwiseConvNode(
+            context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinSum ||
+                 node.builtin_code == kTfLiteBuiltinReduceMin ||
+                 node.builtin_code == kTfLiteBuiltinReduceMax ||
+                 node.builtin_code == kTfLiteBuiltinMean) {
+        TF_LITE_ENSURE_STATUS(
+            DefineReductionNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinStablehloClamp) {
+        TF_LITE_ENSURE_STATUS(DefineStablehloClampNode(
+            context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinQuantize) {
+        TF_LITE_ENSURE_STATUS(
+            DefineQuantizeNode(context, subgraph_, tensor_to_value_id_, node));
+      } else if (node.builtin_code == kTfLiteBuiltinDequantize) {
+        TF_LITE_ENSURE_STATUS(DefineDequantizeNode(context, subgraph_,
+                                                   tensor_to_value_id_, node));
+      } else {
+        TF_LITE_ENSURE_MSG(context, false, "Unsupported op: %d",
+                           node.builtin_code);
+      }
+    }
+
+    ynn_threadpool_t ynn_tp = reinterpret_cast<ynn_threadpool_t>(thread_pool_);
+    TF_LITE_ENSURE_YNN_STATUS(ynn_optimize_subgraph(subgraph_, ynn_tp, 0));
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_create_runtime(subgraph_, ynn_tp, 0, &runtime_));
+
+    return kTfLiteOk;
+  }
+
+  TfLiteStatus Init(TfLiteContext* context,
+                    const TfLiteDelegateParams* params) override {
+    input_tensor_indices_.assign(
+        params->input_tensors->data,
+        params->input_tensors->data + params->input_tensors->size);
+    output_tensor_indices_.assign(
+        params->output_tensors->data,
+        params->output_tensors->data + params->output_tensors->size);
+
+    nodes_info_.clear();
+    nodes_info_.reserve(params->nodes_to_replace->size);
+    for (int i = 0; i < params->nodes_to_replace->size; ++i) {
+      int node_index = params->nodes_to_replace->data[i];
+      TfLiteNode* node;
+      TfLiteRegistration* reg;
+      TF_LITE_ENSURE_STATUS(
+          context->GetNodeAndRegistration(context, node_index, &node, &reg));
+
+      NodeInfo node_info;
+      node_info.node_index = node_index;
+      node_info.builtin_code = reg->builtin_code;
+      node_info.inputs.assign(node->inputs->data,
+                              node->inputs->data + node->inputs->size);
+      node_info.outputs.assign(node->outputs->data,
+                               node->outputs->data + node->outputs->size);
+      node_info.activation = GetFusedActivation(reg, node);
+      nodes_info_.push_back(node_info);
+    }
+
+    return BuildSubgraphAndRuntime(context);
+  }
+
+  TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) override {
+    bool rebuild_required = false;
+    for (size_t i = 0; i < input_tensor_indices_.size(); ++i) {
+      const TfLiteTensor& tensor = context->tensors[input_tensor_indices_[i]];
+      if (options_.static_shape) {
+        if (tensor.dims->size != input_shapes_[i].size() ||
+            !std::equal(tensor.dims->data,
+                        tensor.dims->data + tensor.dims->size,
+                        input_shapes_[i].begin())) {
+          rebuild_required = true;
+          break;
+        }
+      } else {
+        if (tensor.dims->size != input_shapes_[i].size()) {
+          rebuild_required = true;
+          break;
+        }
+      }
+    }
+
+    if (rebuild_required) {
+      TF_LITE_ENSURE_STATUS(BuildSubgraphAndRuntime(context));
+    }
+
+    // Set input shapes in YNNPACK.
+    for (const auto& input : inputs_) {
+      const TfLiteTensor& tensor = context->tensors[input.tensor_index];
+      size_t dims[YNN_MAX_TENSOR_RANK];
+      std::copy_n(tensor.dims->data, tensor.dims->size, dims);
+      TF_LITE_ENSURE_YNN_STATUS(ynn_set_external_value_shape(
+          runtime_, input.val_id, tensor.dims->size, dims));
+    }
+
+    TF_LITE_ENSURE_YNN_STATUS(ynn_reshape_runtime(runtime_));
+
+    // Query output shapes from YNNPACK and resize TFLite tensors.
+    for (const auto& output : outputs_) {
+      TfLiteTensor& tensor = context->tensors[output.tensor_index];
+      size_t rank = YNN_MAX_TENSOR_RANK;
+      size_t dims[YNN_MAX_TENSOR_RANK] = {0};
+      TF_LITE_ENSURE_YNN_STATUS(
+          ynn_get_external_value_shape(runtime_, output.val_id, &rank, dims));
+
+      TfLiteIntArray* output_shape = TfLiteIntArrayCreate(rank);
+      for (size_t d = 0; d < rank; ++d) {
+        constexpr size_t kMaxTfLiteDim = std::numeric_limits<int>::max();
+        TF_LITE_ENSURE(context, dims[d] < kMaxTfLiteDim);
+        output_shape->data[d] = dims[d];
+      }
+      TF_LITE_ENSURE_STATUS(
+          context->ResizeTensor(context, &tensor, output_shape));
+    }
+
+    return kTfLiteOk;
+  }
+
+  TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) override {
+    // Set input buffers.
+    for (const auto& input : inputs_) {
+      TfLiteTensor& tensor = context->tensors[input.tensor_index];
+      TF_LITE_ENSURE_YNN_STATUS(
+          ynn_set_external_value_data(runtime_, input.val_id, tensor.data.raw));
+    }
+
+    // Set output buffers.
+    for (const auto& output : outputs_) {
+      TfLiteTensor& tensor = context->tensors[output.tensor_index];
+      TF_LITE_ENSURE_YNN_STATUS(ynn_set_external_value_data(
+          runtime_, output.val_id, tensor.data.raw));
+    }
+
+    TF_LITE_ENSURE_YNN_STATUS(ynn_invoke_runtime(runtime_));
+    return kTfLiteOk;
+  }
+
+ private:
+  const TfLiteYNNPackDelegateOptions options_;
+  slinky::thread_pool* thread_pool_ = nullptr;
+  ynn_subgraph_t subgraph_;
+  ynn_runtime_t runtime_;
+
+  struct TensorMap {
+    int tensor_index;
+    uint32_t val_id;
+  };
+  std::vector<TensorMap> inputs_;
+  std::vector<TensorMap> outputs_;
+
+  std::vector<NodeInfo> nodes_info_;
+  std::vector<int> input_tensor_indices_;
+  std::vector<int> output_tensor_indices_;
+  std::vector<std::vector<size_t>> input_shapes_;
+  TensorToValueIdMap tensor_to_value_id_;
+};
+
+class YNNPackDelegate : public SimpleDelegateInterface {
+ public:
+  explicit YNNPackDelegate(const TfLiteYNNPackDelegateOptions& options)
+      : options_(options) {
+    if (options_.num_threads > 1) {
+      thread_pool_ =
+          std::make_unique<slinky::thread_pool_impl>(options_.num_threads - 1);
+    }
+  }
+
+  bool IsNodeSupportedByDelegate(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) const override {
+    int builtin_code = registration->builtin_code;
+    if (IsUnaryOp(builtin_code)) {
+      return IsUnaryOpSupported(registration, node, context) == kTfLiteOk;
+    } else if (IsBinaryOp(builtin_code)) {
+      return IsBinaryOpSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinTranspose) {
+      return IsTransposeSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinSlice ||
+               builtin_code == kTfLiteBuiltinStridedSlice) {
+      return IsSliceSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinConcatenation) {
+      return IsConcatenationSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinReshape) {
+      return IsReshapeSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinExpandDims) {
+      return IsExpandDimsSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinPad ||
+               builtin_code == kTfLiteBuiltinPadv2) {
+      return IsPadSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinSplit) {
+      return IsSplitSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinGather) {
+      return IsGatherSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinGatherNd) {
+      return IsGatherNdSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinSpaceToDepth) {
+      return IsSpaceToDepthSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinDepthToSpace) {
+      return IsDepthToSpaceSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinBatchMatmul) {
+      return IsBatchMatMulSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinFullyConnected) {
+      return IsFullyConnectedSupported(registration, node, context) ==
+             kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinConv2d) {
+      return IsConvSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinDepthwiseConv2d) {
+      return IsDepthwiseConvSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinSoftmax ||
+               builtin_code == kTfLiteBuiltinLogSoftmax) {
+      return IsSoftmaxSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinMaxPool2d ||
+               builtin_code == kTfLiteBuiltinAveragePool2d) {
+      return IsPoolingSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinSum ||
+               builtin_code == kTfLiteBuiltinReduceMin ||
+               builtin_code == kTfLiteBuiltinReduceMax ||
+               builtin_code == kTfLiteBuiltinMean) {
+      return IsReductionSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinStablehloClamp) {
+      return IsStablehloClampSupported(registration, node, context) ==
+             kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinQuantize) {
+      return IsQuantizeSupported(registration, node, context) == kTfLiteOk;
+    } else if (builtin_code == kTfLiteBuiltinDequantize) {
+      return IsDequantizeSupported(registration, node, context) == kTfLiteOk;
+    }
+    return false;
+  }
+
+  TfLiteStatus Initialize(TfLiteContext* context) override { return kTfLiteOk; }
+
+  const char* Name() const override {
+    static constexpr char kName[] = "YNNPackDelegate";
+    return kName;
+  }
+
+  std::unique_ptr<SimpleDelegateKernelInterface> CreateDelegateKernelInterface()
+      override {
+    return std::make_unique<YNNPackDelegateKernel>(options_,
+                                                   thread_pool_.get());
+  }
+
+  SimpleDelegateInterface::Options DelegateOptions() const override {
+    return SimpleDelegateInterface::Options();
+  }
+
+ private:
+  const TfLiteYNNPackDelegateOptions options_;
+  std::unique_ptr<slinky::thread_pool_impl> thread_pool_;
+};
+
+}  // namespace ynnpack
+}  // namespace tflite
+
+TfLiteYNNPackDelegateOptions TfLiteYNNPackDelegateOptionsDefault() {
+  TfLiteYNNPackDelegateOptions options = {0};
+  options.num_threads = 1;
+  options.static_shape = false;
+  options.fast_math = false;
+  options.consistent_arithmetic = false;
+  options.no_excess_precision = false;
+  return options;
+}
+
+TfLiteDelegate* TfLiteYNNPackDelegateCreate(
+    const TfLiteYNNPackDelegateOptions* options) {
+  std::unique_ptr<tflite::ynnpack::YNNPackDelegate> ynnpack(
+      new tflite::ynnpack::YNNPackDelegate(
+          options ? *options : TfLiteYNNPackDelegateOptionsDefault()));
+  return tflite::TfLiteDelegateFactory::CreateSimpleDelegate(
+      std::move(ynnpack), kTfLiteDelegateFlagsAllowDynamicTensors);
+}
+
+void TfLiteYNNPackDelegateDelete(TfLiteDelegate* delegate) {
+  tflite::TfLiteDelegateFactory::DeleteSimpleDelegate(delegate);
+}

tensorflow/lite/delegates/ynnpack/ynnpack_delegate.h

@@ -0,0 +1,66 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_YNNPACK_DELEGATE_H_
+#define TENSORFLOW_LITE_DELEGATES_YNNPACK_YNNPACK_DELEGATE_H_
+
+#include <memory>
+
+#include "tensorflow/lite/core/c/common.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+
+typedef struct {
+  // Number of threads to use.
+  int num_threads;
+  // If true, the YNNPACK subgraph will be statically shaped. This can sometimes
+  // improve performance, at the cost of making resizing the graph much more
+  // expensive.
+  bool static_shape;
+  // If true, enable YNN_FLAG_FAST_MATH.
+  bool fast_math;
+  // If true, enable YNN_FLAG_CONSISTENT_ARITHMETIC.
+  bool consistent_arithmetic;
+  // If true, enable YNN_FLAG_NO_EXCESS_PRECISION.
+  bool no_excess_precision;
+} TfLiteYNNPackDelegateOptions;
+
+// Returns a structure with the default delegate options.
+TfLiteYNNPackDelegateOptions TfLiteYNNPackDelegateOptionsDefault();
+
+// Creates a new delegate instance that needs to be destroyed with
+// `TfLiteYNNPackDelegateDelete` when delegate is no longer used by TFLite.
+// When `options` is set to `nullptr`, the default values are used.
+TfLiteDelegate* TfLiteYNNPackDelegateCreate(
+    const TfLiteYNNPackDelegateOptions* options);
+
+// Destroys a delegate created with `TfLiteYNNPackDelegateCreate` call.
+void TfLiteYNNPackDelegateDelete(TfLiteDelegate* delegate);
+
+#ifdef __cplusplus
+}
+#endif  // __cplusplus
+
+// A convenient wrapper that returns C++ std::unique_ptr for automatic memory
+// management.
+inline std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)>
+TfLiteYNNPackDelegateCreateUnique(const TfLiteYNNPackDelegateOptions* options) {
+  return std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)>(
+      TfLiteYNNPackDelegateCreate(options), TfLiteYNNPackDelegateDelete);
+}
+
+#endif  // TENSORFLOW_LITE_DELEGATES_YNNPACK_YNNPACK_DELEGATE_H_