From 22070e377a8a25df2b52dd2ab84af92f2a96e5ec Mon Sep 17 00:00:00 2001
From: Thai Nguyen <thaink@google.com>
Date: Thu, 28 Jan 2021 18:05:32 -0800
Subject: [PATCH] Fix no classes found in selective build android library

The output aar file does not contain any java classes, leads to failure when embedded to Android app.

PiperOrigin-RevId: 354438782
Change-Id: I3f1f0df8f90805142dc28c1d6ccf3c704e7a02bd
---
 tensorflow/lite/build_def.bzl | 16 +++++++++++++---
 tensorflow/lite/java/BUILD    | 23 +++++++++++++----------
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/tensorflow/lite/build_def.bzl b/tensorflow/lite/build_def.bzl
index f166bd602c1..ca283d557c6 100644
--- a/tensorflow/lite/build_def.bzl
+++ b/tensorflow/lite/build_def.bzl
@@ -854,7 +854,9 @@ def tflite_custom_android_library(
         srcs = [],
         deps = [],
         custom_package = "org.tensorflow.lite",
-        visibility = ["//visibility:private"]):
+        visibility = ["//visibility:private"],
+        include_xnnpack_delegate = True,
+        include_nnapi_delegate = True):
     """Generates a tflite Android library, stripping off unused operators.
 
     Note that due to a limitation in the JNI Java wrapper, the compiled TfLite shared binary
@@ -871,9 +873,17 @@ def tflite_custom_android_library(
         custom_package: Name of the Java package. It is required by android_library in case
             the Java source file can't be inferred from the directory where this rule is used.
         visibility: Visibility setting for the generated target. Default to private.
+        include_xnnpack_delegate: Whether to include the XNNPACK delegate or not.
+        include_nnapi_delegate: Whether to include the NNAPI delegate or not.
     """
     tflite_custom_cc_library(name = "%s_cc" % name, models = models, srcs = srcs, deps = deps, visibility = visibility)
 
+    delegate_deps = []
+    if include_nnapi_delegate:
+        delegate_deps.append("//tensorflow/lite/delegates/nnapi/java/src/main/native")
+    if include_xnnpack_delegate:
+        delegate_deps.append("//tensorflow/lite/delegates/xnnpack:xnnpack_delegate")
+
     # JNI wrapper expects a binary file called `libtensorflowlite_jni.so` in java path.
     tflite_jni_binary(
         name = "libtensorflowlite_jni.so",
@@ -882,7 +892,7 @@ def tflite_custom_android_library(
         deps = [
             "//tensorflow/lite/java/src/main/native:native_framework_only",
             ":%s_cc" % name,
-        ],
+        ] + delegate_deps,
     )
 
     native.cc_library(
@@ -893,10 +903,10 @@ def tflite_custom_android_library(
 
     android_library(
         name = name,
+        srcs = ["//tensorflow/lite/java:java_srcs"],
         manifest = "//tensorflow/lite/java:AndroidManifest.xml",
         deps = [
             ":%s_jni" % name,
-            "//tensorflow/lite/java:tensorflowlite_java",
             "@org_checkerframework_qual",
         ],
         custom_package = custom_package,
diff --git a/tensorflow/lite/java/BUILD b/tensorflow/lite/java/BUILD
index 3e915957add..e7755d1e4dc 100644
--- a/tensorflow/lite/java/BUILD
+++ b/tensorflow/lite/java/BUILD
@@ -23,10 +23,13 @@ exports_files([
     "tflite_version_script.lds",
 ])
 
-JAVA_SRCS = glob([
-    "src/main/java/org/tensorflow/lite/*.java",
-    "src/main/java/org/tensorflow/lite/annotations/*.java",
-]) + ["//tensorflow/lite/delegates/nnapi/java/src/main/java/org/tensorflow/lite/nnapi:nnapi_delegate_src"]
+filegroup(
+    name = "java_srcs",
+    srcs = glob([
+        "src/main/java/org/tensorflow/lite/*.java",
+        "src/main/java/org/tensorflow/lite/annotations/*.java",
+    ]) + ["//tensorflow/lite/delegates/nnapi/java/src/main/java/org/tensorflow/lite/nnapi:nnapi_delegate_src"],
+)
 
 # Building tensorflow-lite.aar including 4 variants of .so
 # To build an aar for release, run below command:
@@ -66,7 +69,7 @@ aar_with_jni(
 
 android_library(
     name = "tensorflowlite",
-    srcs = JAVA_SRCS,
+    srcs = [":java_srcs"],
     manifest = "AndroidManifest.xml",
     proguard_specs = ["proguard.flags"],
     deps = [
@@ -111,7 +114,7 @@ android_library(
 
 android_library(
     name = "tensorflowlite_java",
-    srcs = JAVA_SRCS,
+    srcs = [":java_srcs"],
     proguard_specs = ["proguard.flags"],
     deps = [
         "@org_checkerframework_qual",
@@ -120,7 +123,7 @@ android_library(
 
 java_library(
     name = "tensorflowlite_javalib",
-    srcs = JAVA_SRCS,
+    srcs = [":java_srcs"],
     javacopts = JAVACOPTS,
     deps = [
         "@org_checkerframework_qual",
@@ -129,7 +132,7 @@ java_library(
 
 java_library(
     name = "tensorflowlitelib",
-    srcs = JAVA_SRCS,
+    srcs = [":java_srcs"],
     javacopts = JAVACOPTS,
     deps = [
         ":libtensorflowlite_jni.so",
@@ -180,7 +183,7 @@ java_test(
 java_test(
     name = "TensorFlowLiteNoNativeLibTest",
     size = "small",
-    srcs = JAVA_SRCS + ["src/test/java/org/tensorflow/lite/TensorFlowLiteNoNativeLibTest.java"],
+    srcs = [":java_srcs"] + ["src/test/java/org/tensorflow/lite/TensorFlowLiteNoNativeLibTest.java"],
     javacopts = JAVACOPTS,
     test_class = "org.tensorflow.lite.TensorFlowLiteNoNativeLibTest",
     visibility = ["//visibility:private"],
@@ -194,7 +197,7 @@ java_test(
 java_test(
     name = "TensorFlowLiteInvalidNativeLibTest",
     size = "small",
-    srcs = JAVA_SRCS + ["src/test/java/org/tensorflow/lite/TensorFlowLiteInvalidNativeLibTest.java"],
+    srcs = [":java_srcs"] + ["src/test/java/org/tensorflow/lite/TensorFlowLiteInvalidNativeLibTest.java"],
     javacopts = JAVACOPTS,
     test_class = "org.tensorflow.lite.TensorFlowLiteInvalidNativeLibTest",
     visibility = ["//visibility:private"],