chore: import upstream snapshot with attribution

tensorflow/lite/kernels/gradient/BUILD

@@ -0,0 +1,80 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "bcast_grad_args_op",
+    srcs = [
+        "bcast_grad_args.cc",
+    ],
+    hdrs = [
+        "bcast_grad_args.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/kernels/internal:types",
+    ],
+)
+
+cc_library(
+    name = "gradient_ops",
+    srcs = [
+        "gradient_ops.cc",
+    ],
+    hdrs = [
+        "gradient_ops.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":bcast_grad_args_op",
+        "//tensorflow/lite:framework",
+    ],
+)
+
+cc_test(
+    name = "bcast_grad_args_op_test",
+    size = "small",
+    srcs = [
+        "bcast_grad_args_test.cc",
+    ],
+    deps = [
+        ":bcast_grad_args_op",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework_stable",
+        "//tensorflow/lite/kernels:test_main",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
+pybind_extension(
+    name = "pywrap_gradient_ops",
+    srcs = [
+        "gradient_ops_wrapper.cc",
+    ],
+    hdrs = ["gradient_ops.h"],
+    additional_exported_symbols = ["GradientOpsRegisterer"],
+    enable_stub_generation = True,
+    link_in_framework = True,
+    pytype_srcs = [
+        "pywrap_gradient_ops.pyi",
+    ],
+    deps = [
+        ":gradient_ops",
+        "//tensorflow/lite:mutable_op_resolver",
+        "@pybind11",
+    ],
+)

tensorflow/lite/kernels/gradient/bcast_grad_args.cc

@@ -0,0 +1,231 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+// This file implements the TensorFlow Lite's broadcast gradient argument
+// operator.
+
+#include <algorithm>
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace {
+
+static const int kInputOneTensor = 0;
+static const int kInputTwoTensor = 1;
+static const int kOutputOneTensor = 0;
+static const int kOutputTwoTensor = 1;
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  // Check inputs and output.
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  const TfLiteTensor* input1 = GetInput(context, node, kInputOneTensor);
+  TF_LITE_ENSURE(context, input1 != nullptr);
+  const RuntimeShape input1_shape = GetTensorShape(input1);
+  TF_LITE_ENSURE(context,
+                 input1->type == kTfLiteInt32 || input1->type == kTfLiteInt64);
+  TF_LITE_ENSURE_EQ(context, input1_shape.DimensionsCount(), 1);
+
+  const TfLiteTensor* input2 = GetInput(context, node, kInputTwoTensor);
+  TF_LITE_ENSURE(context, input2 != nullptr);
+  const RuntimeShape input2_shape = GetTensorShape(input2);
+  TF_LITE_ENSURE_TYPES_EQ(context, input2->type, input1->type);
+  TF_LITE_ENSURE_EQ(context, input2_shape.DimensionsCount(), 1);
+
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2);
+  TfLiteTensor* output1 = GetOutput(context, node, kOutputOneTensor);
+  TF_LITE_ENSURE(context, output1 != nullptr);
+  TF_LITE_ENSURE_TYPES_EQ(context, output1->type, input1->type);
+  TfLiteTensor* output2 = GetOutput(context, node, kOutputTwoTensor);
+  TF_LITE_ENSURE(context, output2 != nullptr);
+  TF_LITE_ENSURE_TYPES_EQ(context, output2->type, input1->type);
+  SetTensorToDynamic(output1);
+  SetTensorToDynamic(output2);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Invoke(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input1 = GetInput(context, node, kInputOneTensor);
+  TF_LITE_ENSURE(context, input1 != nullptr);
+  const RuntimeShape input1_shape = GetTensorShape(input1);
+
+  const TfLiteTensor* input2 = GetInput(context, node, kInputTwoTensor);
+  TF_LITE_ENSURE(context, input2 != nullptr);
+  const RuntimeShape input2_shape = GetTensorShape(input2);
+
+  TfLiteTensor* output1 = GetOutput(context, node, kOutputOneTensor);
+  TF_LITE_ENSURE(context, output1 != nullptr);
+  TfLiteTensor* output2 = GetOutput(context, node, kOutputTwoTensor);
+  TF_LITE_ENSURE(context, output2 != nullptr);
+
+  std::vector<int64_t> input1_vec;
+  std::vector<int64_t> input2_vec;
+  if (input1->type == kTfLiteInt32) {
+    input1_vec = std::vector<int64_t>(input1->data.i32,
+                                      input1->data.i32 + input1_shape.Dims(0));
+  } else {
+    input1_vec = std::vector<int64_t>(input1->data.i64,
+                                      input1->data.i64 + input1_shape.Dims(0));
+  }
+  if (input2->type == kTfLiteInt32) {
+    input2_vec = std::vector<int64_t>(input2->data.i32,
+                                      input2->data.i32 + input2_shape.Dims(0));
+  } else {
+    input2_vec = std::vector<int64_t>(input2->data.i64,
+                                      input2->data.i64 + input2_shape.Dims(0));
+  }
+
+  if (input1_vec == input2_vec) {
+    // All equals.
+    TfLiteIntArray* output1_shape = TfLiteIntArrayCreate(1);
+    output1_shape->data[0] = 0;
+    TF_LITE_ENSURE_OK(context,
+                      context->ResizeTensor(context, output1, output1_shape));
+
+    TfLiteIntArray* output2_shape = TfLiteIntArrayCreate(1);
+    output2_shape->data[0] = 0;
+    TF_LITE_ENSURE_OK(context,
+                      context->ResizeTensor(context, output2, output2_shape));
+    return kTfLiteOk;
+  }
+
+  size_t largest_rank = std::max(input1_vec.size(), input2_vec.size());
+
+  // Reverse all the shapes for convenience
+  // After the reverse, 0-th is the inner-most dimension.
+  std::vector<int64_t> copy[2];
+  copy[0] = std::vector<int64_t>(input1_vec.rbegin(), input1_vec.rend());
+  copy[1] = std::vector<int64_t>(input2_vec.rbegin(), input2_vec.rend());
+
+  // 1-extend and align all vectors.
+  for (int i = 0; i < 2; ++i) {
+    if (copy[i].size() < largest_rank) {
+      copy[i].resize(largest_rank, 1);
+    }
+  }
+  // Going through each dimension starting from the inner-most
+  // dimension, compares dimension of x and y. They are compatible if
+  // they are equal or either is 1.
+
+  // indices of j-th component of each input.
+  std::array<bool, 2> prev_is_one = {false, false};
+  std::array<bool, 2> current_is_one = {false, false};
+  bool set_one = false;
+  // indices of gradient reduction of each input.
+  std::vector<int64_t> grad_reduce_idx[2];
+
+  for (int j = 0; j < largest_rank; ++j) {
+    int output_dim = -1;
+    int output_dim_set = false;
+    bool none_is_one = true;
+    // Find which indices are 1.
+    for (int i = 0; i < 2; ++i) {
+      // Keep track of which indices are 1.
+      if (copy[i][j] == 1) {
+        current_is_one[i] = true;
+        none_is_one = false;
+      } else {
+        current_is_one[i] = false;
+        if (!output_dim_set || copy[i][j] == output_dim) {
+          output_dim = copy[i][j];
+          output_dim_set = true;
+        } else {
+          // Not broadcastable shapes.
+          return kTfLiteError;
+        }
+      }
+    }
+    // All dimensions are 1.
+    if (!output_dim_set) {
+      for (int i = 0; i < 2; ++i) {
+        grad_reduce_idx[i].push_back(largest_rank - 1 - j);
+      }
+      continue;
+    } else if (current_is_one == prev_is_one && set_one) {
+      // It is a run of the same broadcasting case as last time.
+      // We can reshape the input so that fewer dimensions
+      // are involved in the intermediate computation.
+      for (int i = 0; i < 2; ++i) {
+        if (current_is_one[i] && !none_is_one) {
+          grad_reduce_idx[i].push_back(largest_rank - 1 - j);
+        }
+      }
+    } else {
+      for (int i = 0; i < 2; ++i) {
+        if (current_is_one[i] && !none_is_one) {
+          grad_reduce_idx[i].push_back(largest_rank - 1 - j);
+        }
+      }
+    }
+    set_one = true;
+    for (int i = 0; i < 2; ++i) {
+      prev_is_one[i] = current_is_one[i];
+    }
+  }
+  for (int i = 0; i < 2; ++i) {
+    std::reverse(grad_reduce_idx[i].begin(), grad_reduce_idx[i].end());
+  }
+  TfLiteIntArray* output1_shape = TfLiteIntArrayCreate(1);
+  output1_shape->data[0] = grad_reduce_idx[0].size();
+  TF_LITE_ENSURE_OK(context,
+                    context->ResizeTensor(context, output1, output1_shape));
+  if (output1->type == kTfLiteInt32) {
+    for (int i = 0; i < grad_reduce_idx[0].size(); ++i) {
+      output1->data.i32[i] = grad_reduce_idx[0][i];
+    }
+  } else if (output1->type == kTfLiteInt64) {
+    for (int i = 0; i < grad_reduce_idx[0].size(); ++i) {
+      output1->data.i64[i] = grad_reduce_idx[0][i];
+    }
+  }
+
+  TfLiteIntArray* output2_shape = TfLiteIntArrayCreate(1);
+  output2_shape->data[0] = grad_reduce_idx[1].size();
+  TF_LITE_ENSURE_OK(context,
+                    context->ResizeTensor(context, output2, output2_shape));
+  if (output2->type == kTfLiteInt32) {
+    for (int i = 0; i < grad_reduce_idx[1].size(); ++i) {
+      output2->data.i32[i] = grad_reduce_idx[1][i];
+    }
+  } else if (output2->type == kTfLiteInt64) {
+    for (int i = 0; i < grad_reduce_idx[1].size(); ++i) {
+      output2->data.i64[i] = grad_reduce_idx[1][i];
+    }
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace
+
+TfLiteRegistration* Register_BROADCAST_GRADIENT_ARGS() {
+  static TfLiteRegistration reg = {/*init=*/nullptr,
+                                   /*free=*/nullptr,
+                                   /*prepare=*/Prepare,
+                                   /*invoke=*/Invoke};
+  return &reg;
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/gradient/bcast_grad_args.h

@@ -0,0 +1,34 @@
+/* Copyright 2021 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_KERNELS_GRADIENT_BCAST_GRAD_ARGS_H_
+#define TENSORFLOW_LITE_KERNELS_GRADIENT_BCAST_GRAD_ARGS_H_
+
+// This file declares the TensorFlow Lite's broadcast gradient argument custom
+// operator.
+
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+TfLiteRegistration* Register_BROADCAST_GRADIENT_ARGS();
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_GRADIENT_BCAST_GRAD_ARGS_H_

tensorflow/lite/kernels/gradient/bcast_grad_args_test.cc

@@ -0,0 +1,242 @@
+/* Copyright 2021 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/kernels/gradient/bcast_grad_args.h"
+
+#include <cstdint>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+namespace {
+
+using testing::ElementsAreArray;
+
+class BcastGradArgsInt32OpModel : public SingleOpModel {
+ public:
+  BcastGradArgsInt32OpModel(const TensorData& input1, const TensorData& input2,
+                            const TensorData& output1,
+                            const TensorData& output2) {
+    input1_ = AddInput(input1);
+    input2_ = AddInput(input2);
+    output1_ = AddOutput(output1);
+    output2_ = AddOutput(output2);
+
+    std::vector<uint8_t> custom_option;
+    SetCustomOp("BroadcastGradientArgs", custom_option,
+                Register_BROADCAST_GRADIENT_ARGS);
+    BuildInterpreter({GetShape(input1_), GetShape(input2_)});
+  }
+
+  void SetInput1(const std::vector<int>& data) {
+    PopulateTensor(input1_, data);
+  }
+  void SetInput2(const std::vector<int>& data) {
+    PopulateTensor(input2_, data);
+  }
+
+  std::vector<int> GetOutput1() { return ExtractVector<int>(output1_); }
+  std::vector<int> GetOutput1Shape() { return GetTensorShape(output1_); }
+  std::vector<int> GetOutput2() { return ExtractVector<int>(output2_); }
+  std::vector<int> GetOutput2Shape() { return GetTensorShape(output2_); }
+
+ protected:
+  int input1_;
+  int input2_;
+  int output1_;
+  int output2_;
+};
+
+TEST(BcastGradArgsInt32OpModel, AllEqualsInt32DTypes) {
+  BcastGradArgsInt32OpModel model(
+      /*input1=*/{TensorType_INT32, {4}},
+      /*input2=*/{TensorType_INT32, {4}},
+      /*output1=*/{TensorType_INT32, {}},
+      /*output2=*/{TensorType_INT32, {}});
+  model.SetInput1({3, 1, 2, 3});
+  model.SetInput2({3, 1, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1().size(), 0);
+  EXPECT_THAT(model.GetOutput2().size(), 0);
+}
+
+TEST(BcastGradArgsInt32OpModel, BroadcastableDimAtInput1Int32DTypes) {
+  BcastGradArgsInt32OpModel model(
+      /*input1=*/{TensorType_INT32, {4}},
+      /*input2=*/{TensorType_INT32, {4}},
+      /*output1=*/{TensorType_INT32, {}},
+      /*output2=*/{TensorType_INT32, {}});
+  model.SetInput1({3, 4, 1, 3});
+  model.SetInput2({3, 4, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1(), ElementsAreArray({2}));
+  EXPECT_THAT(model.GetOutput2().size(), 0);
+}
+
+TEST(BcastGradArgsInt32OpModel, BroadcastableDimAtInput2Int32DTypes) {
+  BcastGradArgsInt32OpModel model(
+      /*input1=*/{TensorType_INT32, {4}},
+      /*input2=*/{TensorType_INT32, {4}},
+      /*output1=*/{TensorType_INT32, {}},
+      /*output2=*/{TensorType_INT32, {}});
+  model.SetInput1({3, 4, 2, 3});
+  model.SetInput2({3, 1, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1().size(), 0);
+  EXPECT_THAT(model.GetOutput2(), ElementsAreArray({1}));
+}
+
+TEST(BcastGradArgsInt32OpModel, DifferentInputSizesInt32DTypes) {
+  BcastGradArgsInt32OpModel model(
+      /*input1=*/{TensorType_INT32, {4}},
+      /*input2=*/{TensorType_INT32, {3}},
+      /*output1=*/{TensorType_INT32, {}},
+      /*output2=*/{TensorType_INT32, {}});
+  model.SetInput1({3, 4, 2, 3});
+  model.SetInput2({4, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1().size(), 0);
+  EXPECT_THAT(model.GetOutput2(), ElementsAreArray({0}));
+}
+
+TEST(BcastGradArgsInt32OpModel, NonBroadcastableDimsInt32DTypes) {
+  BcastGradArgsInt32OpModel model(
+      /*input1=*/{TensorType_INT32, {4}},
+      /*input2=*/{TensorType_INT32, {4}},
+      /*output1=*/{TensorType_INT32, {}},
+      /*output2=*/{TensorType_INT32, {}});
+  model.SetInput1({3, 4, 2, 3});
+  model.SetInput2({9, 9, 9, 9});
+  EXPECT_THAT(model.Invoke(), kTfLiteError);
+}
+
+class BcastGradArgsInt64OpModel : public SingleOpModel {
+ public:
+  BcastGradArgsInt64OpModel(const TensorData& input1, const TensorData& input2,
+                            const TensorData& output1,
+                            const TensorData& output2) {
+    input1_ = AddInput(input1);
+    input2_ = AddInput(input2);
+    output1_ = AddOutput(output1);
+    output2_ = AddOutput(output2);
+
+    std::vector<uint8_t> custom_option;
+    SetCustomOp("BroadcastGradientArgs", custom_option,
+                Register_BROADCAST_GRADIENT_ARGS);
+    BuildInterpreter({GetShape(input1_), GetShape(input2_)});
+  }
+
+  void SetInput1(const std::vector<int64_t>& data) {
+    PopulateTensor(input1_, data);
+  }
+  void SetInput2(const std::vector<int64_t>& data) {
+    PopulateTensor(input2_, data);
+  }
+
+  std::vector<int64_t> GetOutput1() { return ExtractVector<int64_t>(output1_); }
+  std::vector<int> GetOutput1Shape() { return GetTensorShape(output1_); }
+  std::vector<int64_t> GetOutput2() { return ExtractVector<int64_t>(output2_); }
+  std::vector<int> GetOutput2Shape() { return GetTensorShape(output2_); }
+
+ protected:
+  int input1_;
+  int input2_;
+  int output1_;
+  int output2_;
+};
+
+TEST(BcastGradArgsInt32OpModel, AllEqualsInt64DTypes) {
+  BcastGradArgsInt64OpModel model(
+      /*input1=*/{TensorType_INT64, {4}},
+      /*input2=*/{TensorType_INT64, {4}},
+      /*output1=*/{TensorType_INT64, {}},
+      /*output2=*/{TensorType_INT64, {}});
+  model.SetInput1({3, 1, 2, 3});
+  model.SetInput2({3, 1, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1().size(), 0);
+  EXPECT_THAT(model.GetOutput2().size(), 0);
+}
+
+TEST(BcastGradArgsInt32OpModel, BroadcastableDimAtInput1Int64DTypes) {
+  BcastGradArgsInt64OpModel model(
+      /*input1=*/{TensorType_INT64, {4}},
+      /*input2=*/{TensorType_INT64, {4}},
+      /*output1=*/{TensorType_INT64, {}},
+      /*output2=*/{TensorType_INT64, {}});
+  model.SetInput1({3, 4, 1, 3});
+  model.SetInput2({3, 4, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1(), ElementsAreArray({2}));
+  EXPECT_THAT(model.GetOutput2().size(), 0);
+}
+
+TEST(BcastGradArgsInt32OpModel, BroadcastableDimAtInput2Int64DTypes) {
+  BcastGradArgsInt64OpModel model(
+      /*input1=*/{TensorType_INT64, {4}},
+      /*input2=*/{TensorType_INT64, {4}},
+      /*output1=*/{TensorType_INT64, {}},
+      /*output2=*/{TensorType_INT64, {}});
+  model.SetInput1({3, 4, 2, 3});
+  model.SetInput2({3, 1, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1().size(), 0);
+  EXPECT_THAT(model.GetOutput2(), ElementsAreArray({1}));
+}
+
+TEST(BcastGradArgsInt32OpModel, DifferentInputSizesInt64DTypes) {
+  BcastGradArgsInt64OpModel model(
+      /*input1=*/{TensorType_INT64, {4}},
+      /*input2=*/{TensorType_INT64, {3}},
+      /*output1=*/{TensorType_INT64, {}},
+      /*output2=*/{TensorType_INT64, {}});
+  model.SetInput1({3, 4, 2, 3});
+  model.SetInput2({4, 2, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput1().size(), 0);
+  EXPECT_THAT(model.GetOutput2(), ElementsAreArray({0}));
+}
+
+TEST(BcastGradArgsInt32OpModel, NonBroadcastableDimsInt64DTypes) {
+  BcastGradArgsInt64OpModel model(
+      /*input1=*/{TensorType_INT64, {4}},
+      /*input2=*/{TensorType_INT64, {4}},
+      /*output1=*/{TensorType_INT64, {}},
+      /*output2=*/{TensorType_INT64, {}});
+  model.SetInput1({3, 4, 2, 3});
+  model.SetInput2({9, 9, 9, 9});
+  EXPECT_THAT(model.Invoke(), kTfLiteError);
+}
+
+}  // namespace
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/gradient/gradient_ops.cc

@@ -0,0 +1,32 @@
+/* Copyright 2021 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/kernels/gradient/gradient_ops.h"
+
+#include "tensorflow/lite/kernels/gradient/bcast_grad_args.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+extern "C" void AddGradientOps(::tflite::MutableOpResolver* resolver) {
+  resolver->AddCustom("BroadcastGradientArgs",
+                      tflite::ops::custom::Register_BROADCAST_GRADIENT_ARGS());
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/gradient/gradient_ops.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 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_KERNELS_GRADIENT_GRADIENT_OPS_H_
+#define TENSORFLOW_LITE_KERNELS_GRADIENT_GRADIENT_OPS_H_
+
+// This file declares the TensorFlow Lite's gradient custom operators.
+
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+extern "C" void AddGradientOps(::tflite::MutableOpResolver* resolver);
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_GRADIENT_GRADIENT_OPS_H_

tensorflow/lite/kernels/gradient/gradient_ops_wrapper.cc

@@ -0,0 +1,37 @@
+/* Copyright 2021 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 <cstdint>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "tensorflow/lite/kernels/gradient/gradient_ops.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+PYBIND11_MODULE(pywrap_gradient_ops, m) {
+  m.doc() = R"pbdoc(
+    pywrap_gradient_ops
+    -----
+  )pbdoc";
+  m.def(
+      "GradientOpsRegisterer",
+      [](uintptr_t resolver) {
+        tflite::ops::custom::AddGradientOps(
+            reinterpret_cast<tflite::MutableOpResolver*>(resolver));
+      },
+      R"pbdoc(
+        Gradient op registerer function with the correct signature. Registers
+        Gradient custom ops.
+      )pbdoc");
+}

tensorflow/lite/kernels/gradient/pywrap_gradient_ops.pyi

@@ -0,0 +1,16 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+def GradientOpsRegisterer(arg0: int) -> None: ...