Skip to content

Commit

Permalink
Port BATCH_MATMUL (#2147)
Browse files Browse the repository at this point in the history
@tensorflow/micro

Port BATCH_MATMUL kernel from TfLite along with unit tests.

[b/291110349
](https://issuetracker.google.com/issues/291110349)

bug=fixes #2146
  • Loading branch information
ddavis-2015 authored Sep 6, 2023
1 parent 0cc7719 commit 5d542d7
Show file tree
Hide file tree
Showing 11 changed files with 1,322 additions and 7 deletions.
15 changes: 15 additions & 0 deletions tensorflow/lite/micro/kernels/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@ tflm_kernel_cc_library(
"add_n.cc",
"arg_min_max.cc",
"assign_variable.cc",
"batch_matmul.cc",
"batch_to_space_nd.cc",
"broadcast_args.cc",
"broadcast_to.cc",
Expand Down Expand Up @@ -425,6 +426,20 @@ cc_test(
],
)

cc_test(
name = "batch_matmul_test",
srcs = [
"batch_matmul_test.cc",
],
deps = [
":kernel_runner",
"//tensorflow/lite/c:common",
"//tensorflow/lite/micro:op_resolvers",
"//tensorflow/lite/micro:test_helpers",
"//tensorflow/lite/micro/testing:micro_test",
],
)

cc_test(
name = "batch_to_space_nd_test",
srcs = [
Expand Down
3 changes: 2 additions & 1 deletion tensorflow/lite/micro/kernels/Makefile.inc
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
# 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.
Expand Down Expand Up @@ -114,6 +114,7 @@ $(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/activations_test.cc \
$(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/add_test.cc \
$(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/add_n_test.cc \
$(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/arg_min_max_test.cc \
$(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/batch_matmul_test.cc \
$(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/batch_to_space_nd_test.cc \
$(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/broadcast_args_test.cc \
$(TENSORFLOW_ROOT)tensorflow/lite/micro/kernels/broadcast_to_test.cc \
Expand Down
Loading

0 comments on commit 5d542d7

Please sign in to comment.