Merge "Add lingvo assert ops to the tff_server container." into main

WORKSPACE

@@ -110,6 +110,14 @@ load("@org_tensorflow//tensorflow:workspace0.bzl", "tf_workspace0")
 
 tf_workspace0()
 
+http_archive(
+    name = "lingvo",
+    patches = ["//third_party/lingvo:lingvo.bzl.patch"],
+    sha256 = "6baf671a81ed747f2580eb4f549044093c48173cd88b392510fe6ddf5dce2ba2",
+    strip_prefix = "lingvo-ccfa97995bea99a3c0bb47b7b0b8e34a757ecf39",
+    url = "https://github.com/tensorflow/lingvo/archive/ccfa97995bea99a3c0bb47b7b0b8e34a757ecf39.tar.gz",
+)
+
 load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies", "rules_proto_toolchains")
 
 rules_proto_dependencies()

containers/tff_server/BUILD

@@ -38,6 +38,7 @@ cc_library(
         "@federated-compute//fcp/protos/confidentialcompute:confidential_transform_cc_grpc",
         "@federated-compute//fcp/protos/confidentialcompute:confidential_transform_cc_proto",
         "@federated-compute//fcp/protos/confidentialcompute:tff_config_cc_proto",
+        "@lingvo//lingvo/core/ops:assert_kernels",
         "@oak//proto/containers:orchestrator_crypto_cc_grpc",
         "@org_tensorflow//tensorflow/compiler/jit:xla_cpu_jit",
         "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor",

third_party/lingvo/lingvo.bzl.patch

@@ -0,0 +1,42 @@
+# Patch custom_kernel_library to use tf_kernel_library directly. The lingvo
+# implementation replies on a tf_includes repo which is created inside lingvo
+# based on TF 2.13. This patch potentially can be removed once lingvo moves to
+# newer version of TensorFlow.
+diff --git lingvo/lingvo.bzl lingvo/lingvo.bzl
+index 251e25f88..83fd7f2fd 100644
+--- lingvo/lingvo.bzl
++++ lingvo/lingvo.bzl
+@@ -1,4 +1,5 @@
+ """Implements custom rules for Lingvo."""
++load("@org_tensorflow//tensorflow:tensorflow.bzl", "tf_kernel_library")
+
+ def tf_copts():
+     #  "-Wno-sign-compare", "-mavx" removed for compat with aarch64
+@@ -57,15 +58,18 @@ def lingvo_cc_test(name, srcs, deps = [], **kwargs):
+ lingvo_py_binary = native.py_binary
+
+ def custom_kernel_library(name, op_def_lib, srcs, hdrs = [], deps = []):
+-    native.cc_library(
+-        name = name,
+-        srcs = srcs,
+-        hdrs = hdrs,
+-        copts = tf_copts(),
+-        deps = [
+-            "@tensorflow_includes//:includes",
+-        ] + deps + op_def_lib,
+-        alwayslink = 1,
++    tf_deps = [
++       "@org_tensorflow//tensorflow/core:core_cpu",
++       "@org_tensorflow//tensorflow/core:framework",
++       "@org_tensorflow//tensorflow/core:lib",
++    ]
++    tf_kernel_library(
++       name = name,
++       srcs = srcs,
++       hdrs = hdrs,
++       copts = tf_copts(),
++       deps = deps + tf_deps,
++       alwayslink = 1,
+     )
+
+ def gen_op_cclib(name, srcs, deps = [], nonportable_deps = []):