Update FCP to use the new location of the aggregation code.upstream-main

PiperOrigin-RevId: 630470255

10 files changed, 89 insertions, 72 deletions

diff --git a/WORKSPACE b/WORKSPACE
index e543ef3..b1ab0be 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -323,3 +323,10 @@ http_archive(
   sha256 = "6e6e1cd7a819695caae408f4fa938129ab7a86e83fe2410137c85e50131abbe0",
   strip_prefix = "differential-privacy-3.0.0",
 )
+
+http_archive(
+    name = "org_tensorflow_federated",
+    url = "https://github.com/tensorflow/federated/archive/93ffd03340d021a336994ace52ea6919b1821ff5.tar.gz",
+    sha256 = "247295d40ab2c78c22dec44b36a427fa143bb5a343a3a46f91fcddad5ac2ee07",
+    strip_prefix = "federated-93ffd03340d021a336994ace52ea6919b1821ff5",
+)
diff --git a/fcp/client/BUILD b/fcp/client/BUILD
index 4c08cfd..608023c 100644
--- a/fcp/client/BUILD
+++ b/fcp/client/BUILD
@@ -543,12 +543,12 @@ cc_library(
     hdrs = ["converters.h"],
     copts = FCP_COPTS,
     deps = [
-        "//fcp/aggregation/core:tensor",
-        "//fcp/aggregation/core:tensor_cc_proto",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@com_google_protobuf//:protobuf",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor_cc_proto",
     ],
 )
 
@@ -558,11 +558,11 @@ cc_test(
     deps = [
         ":converters",
         ":example_query_result_cc_proto",
-        "//fcp/aggregation/core:tensor",
-        "//fcp/aggregation/core:tensor_cc_proto",
-        "//fcp/aggregation/testing",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor_cc_proto",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/testing",
     ],
 )
 
diff --git a/fcp/client/converters.cc b/fcp/client/converters.cc
index d902820..7e95492 100644
--- a/fcp/client/converters.cc
+++ b/fcp/client/converters.cc
@@ -21,19 +21,19 @@
 
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "fcp/aggregation/core/tensor.h"
-#include "fcp/aggregation/core/tensor.pb.h"
-#include "fcp/aggregation/core/tensor_data.h"
-#include "fcp/aggregation/core/tensor_shape.h"
 #include "google/protobuf/repeated_ptr_field.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.pb.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_data.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_shape.h"
 
 namespace fcp::client {
 
 using ::absl::StatusOr;
 using ::absl::string_view;
-using ::fcp::aggregation::Tensor;
-using ::fcp::aggregation::TensorData;
-using ::fcp::aggregation::TensorShape;
+using tensorflow_federated::aggregation::Tensor;
+using tensorflow_federated::aggregation::TensorData;
+using tensorflow_federated::aggregation::TensorShape;
 
 // Similar to NumericTensorDataAdapter but performs additional conversion
 // of the original value to string_view while keeping the reference to the
@@ -57,7 +57,8 @@ class StringTensorDataAdapter : public TensorData {
 StatusOr<Tensor> ConvertStringTensor(
     TensorShape tensor_shape,
     const ::google::protobuf::RepeatedPtrField<std::string>& value) {
-  return Tensor::Create(aggregation::DT_STRING, tensor_shape,
+  return Tensor::Create(tensorflow_federated::aggregation::DT_STRING,
+                        tensor_shape,
                         std::make_unique<StringTensorDataAdapter>(value));
 }
 
diff --git a/fcp/client/converters.h b/fcp/client/converters.h
index 6a02db9..f98b3d4 100644
--- a/fcp/client/converters.h
+++ b/fcp/client/converters.h
@@ -22,18 +22,19 @@
 
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
-#include "fcp/aggregation/core/tensor.h"
-#include "fcp/aggregation/core/tensor.pb.h"
-#include "fcp/aggregation/core/tensor_data.h"
-#include "fcp/aggregation/core/tensor_shape.h"
 #include "google/protobuf/repeated_ptr_field.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.pb.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_data.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_shape.h"
 
 namespace fcp::client {
 
 // A primitive TensorData implementation that holds reference to the repeated
 // value field in ExampleQueryResult values.
 template <typename T>
-class NumericTensorDataAdapter : public aggregation::TensorData {
+class NumericTensorDataAdapter
+    : public tensorflow_federated::aggregation::TensorData {
  public:
   explicit NumericTensorDataAdapter(absl::Span<const T> value)
       : value_(value) {}
@@ -48,18 +49,19 @@ class NumericTensorDataAdapter : public aggregation::TensorData {
 // Converts repeated numeric field of ExampleQueryResult Values to Aggregation
 // Tensor.
 template <typename Numeric>
-absl::StatusOr<aggregation::Tensor> ConvertNumericTensor(
-    aggregation::DataType dtype, aggregation::TensorShape tensor_shape,
+absl::StatusOr<tensorflow_federated::aggregation::Tensor> ConvertNumericTensor(
+    tensorflow_federated::aggregation::DataType dtype,
+    tensorflow_federated::aggregation::TensorShape tensor_shape,
     absl::Span<const Numeric> value) {
-  return aggregation::Tensor::Create(
+  return tensorflow_federated::aggregation::Tensor::Create(
       dtype, tensor_shape,
       std::make_unique<NumericTensorDataAdapter<Numeric>>(value));
 }
 
 // Converts repeated string field of ExampleQueryResult Values to Aggregation
 // Tensor.
-absl::StatusOr<aggregation::Tensor> ConvertStringTensor(
-    aggregation::TensorShape tensor_shape,
+absl::StatusOr<tensorflow_federated::aggregation::Tensor> ConvertStringTensor(
+    tensorflow_federated::aggregation::TensorShape tensor_shape,
     const ::google::protobuf::RepeatedPtrField<std::string>& value);
 }  // namespace fcp::client
 
diff --git a/fcp/client/converters_test.cc b/fcp/client/converters_test.cc
index fb89a06..9366913 100644
--- a/fcp/client/converters_test.cc
+++ b/fcp/client/converters_test.cc
@@ -17,10 +17,10 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "absl/strings/string_view.h"
-#include "fcp/aggregation/core/tensor.pb.h"
-#include "fcp/aggregation/core/tensor_shape.h"
-#include "fcp/aggregation/testing/testing.h"
 #include "fcp/client/example_query_result.pb.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.pb.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_shape.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/testing/testing.h"
 
 namespace fcp::client {
 namespace {
@@ -31,10 +31,12 @@ TEST(ConvertersTest, ConvertNumericTensor_Success) {
   float_values.mutable_float_values()->add_value(0.42f);
   float_values.mutable_float_values()->add_value(0.33f);
 
-  EXPECT_THAT(*ConvertNumericTensor<float>(aggregation::DT_FLOAT,
-                                           aggregation::TensorShape({3}),
-                                           float_values.float_values().value()),
-              aggregation::IsTensor<float>({3}, {0.24f, 0.42f, 0.33f}));
+  EXPECT_THAT(*ConvertNumericTensor<float>(
+                  tensorflow_federated::aggregation::DT_FLOAT,
+                  tensorflow_federated::aggregation::TensorShape({3}),
+                  float_values.float_values().value()),
+              tensorflow_federated::aggregation::IsTensor<float>(
+                  {3}, {0.24f, 0.42f, 0.33f}));
 }
 
 TEST(ConvertersTest, ConvertStringTensor_Success) {
@@ -42,10 +44,11 @@ TEST(ConvertersTest, ConvertStringTensor_Success) {
   string_values.mutable_string_values()->add_value("string_value1");
   string_values.mutable_string_values()->add_value("string_value2");
 
-  EXPECT_THAT(*ConvertStringTensor(aggregation::TensorShape({2}),
-                                   string_values.string_values().value()),
-              aggregation::IsTensor<absl::string_view>(
-                  {2}, {"string_value1", "string_value2"}));
+  EXPECT_THAT(
+      *ConvertStringTensor(tensorflow_federated::aggregation::TensorShape({2}),
+                           string_values.string_values().value()),
+      tensorflow_federated::aggregation::IsTensor<absl::string_view>(
+          {2}, {"string_value1", "string_value2"}));
 }
 }  // namespace
 }  // namespace fcp::client
diff --git a/fcp/client/engine/BUILD b/fcp/client/engine/BUILD
index 09ff427..ed48702 100644
--- a/fcp/client/engine/BUILD
+++ b/fcp/client/engine/BUILD
@@ -70,10 +70,6 @@ cc_library(
         ":common",
         ":example_iterator_factory",
         ":plan_engine_helpers",
-        "//fcp/aggregation/core:tensor",
-        "//fcp/aggregation/core:tensor_cc_proto",
-        "//fcp/aggregation/protocol:checkpoint_builder",
-        "//fcp/aggregation/protocol:federated_compute_checkpoint_builder",
         "//fcp/base",
         "//fcp/client:converters",
         "//fcp/client:example_iterator_query_recorder",
@@ -87,6 +83,10 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core/platform:tstring",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor_cc_proto",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/protocol:checkpoint_builder",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/protocol:federated_compute_checkpoint_builder",
     ],
 )
 
@@ -97,10 +97,6 @@ cc_test(
         ":common",
         ":example_iterator_factory",
         ":example_query_plan_engine",
-        "//fcp/aggregation/core:tensor",
-        "//fcp/aggregation/core:tensor_cc_proto",
-        "//fcp/aggregation/protocol:checkpoint_header",
-        "//fcp/aggregation/testing:test_data",
         "//fcp/client:client_runner",
         "//fcp/client:example_query_result_cc_proto",
         "//fcp/client:simple_task_environment",
@@ -122,6 +118,10 @@ cc_test(
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core/platform:tstring",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/core:tensor_cc_proto",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/protocol:checkpoint_header",
+        "@org_tensorflow_federated//tensorflow_federated/cc/core/impl/aggregation/testing:test_data",
     ],
 )
 
diff --git a/fcp/client/engine/example_query_plan_engine.cc b/fcp/client/engine/example_query_plan_engine.cc
index bb13ec1..bbbe1b2 100644
--- a/fcp/client/engine/example_query_plan_engine.cc
+++ b/fcp/client/engine/example_query_plan_engine.cc
@@ -27,11 +27,6 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "fcp/aggregation/core/tensor.h"
-#include "fcp/aggregation/core/tensor.pb.h"
-#include "fcp/aggregation/core/tensor_shape.h"
-#include "fcp/aggregation/protocol/checkpoint_builder.h"
-#include "fcp/aggregation/protocol/federated_compute_checkpoint_builder.h"
 #include "fcp/base/monitoring.h"
 #include "fcp/client/converters.h"
 #include "fcp/client/engine/common.h"
@@ -47,6 +42,11 @@
 #include "tensorflow/core/framework/tensor_slice.h"
 #include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/util/tensor_slice_writer.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.pb.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_shape.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_builder.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/protocol/federated_compute_checkpoint_builder.h"
 
 namespace fcp {
 namespace client {
@@ -54,14 +54,14 @@ namespace engine {
 
 namespace tf = ::tensorflow;
 
-using ::fcp::aggregation::CheckpointBuilder;
-using ::fcp::aggregation::Tensor;
-using ::fcp::aggregation::TensorShape;
 using ::fcp::client::ExampleQueryResult;
 using ::fcp::client::engine::PlanResult;
 using ::fcp::client::opstats::OpStatsLogger;
 using ::google::internal::federated::plan::ExampleQuerySpec;
 using ::google::internal::federated::plan::ExampleSelector;
+using tensorflow_federated::aggregation::CheckpointBuilder;
+using tensorflow_federated::aggregation::Tensor;
+using tensorflow_federated::aggregation::TensorShape;
 
 namespace {
 
@@ -217,7 +217,7 @@ absl::Status GenerateAggregationTensors(
             output_vector_spec, ExampleQuerySpec::OutputVectorSpec::INT32));
         FCP_ASSIGN_OR_RETURN(
             tensor, ConvertNumericTensor<int32_t>(
-                        aggregation::DT_INT32,
+                        tensorflow_federated::aggregation::DT_INT32,
                         TensorShape({values.int32_values().value_size()}),
                         values.int32_values().value()));
 
@@ -226,7 +226,7 @@ absl::Status GenerateAggregationTensors(
             output_vector_spec, ExampleQuerySpec::OutputVectorSpec::INT64));
         FCP_ASSIGN_OR_RETURN(
             tensor, ConvertNumericTensor<int64_t>(
-                        aggregation::DT_INT64,
+                        tensorflow_federated::aggregation::DT_INT64,
                         TensorShape({values.int64_values().value_size()}),
                         values.int64_values().value()));
       } else if (values.has_string_values()) {
@@ -244,7 +244,7 @@ absl::Status GenerateAggregationTensors(
             output_vector_spec, ExampleQuerySpec::OutputVectorSpec::FLOAT));
         FCP_ASSIGN_OR_RETURN(
             tensor, ConvertNumericTensor<float>(
-                        aggregation::DT_FLOAT,
+                        tensorflow_federated::aggregation::DT_FLOAT,
                         TensorShape({values.float_values().value_size()}),
                         values.float_values().value()));
       } else if (values.has_double_values()) {
@@ -252,7 +252,7 @@ absl::Status GenerateAggregationTensors(
             output_vector_spec, ExampleQuerySpec::OutputVectorSpec::DOUBLE));
         FCP_ASSIGN_OR_RETURN(
             tensor, ConvertNumericTensor<double>(
-                        aggregation::DT_DOUBLE,
+                        tensorflow_federated::aggregation::DT_DOUBLE,
                         TensorShape({values.double_values().value_size()}),
                         values.double_values().value()));
       } else if (values.has_bytes_values()) {
diff --git a/fcp/client/engine/example_query_plan_engine.h b/fcp/client/engine/example_query_plan_engine.h
index d1b4bf3..8b080b7 100644
--- a/fcp/client/engine/example_query_plan_engine.h
+++ b/fcp/client/engine/example_query_plan_engine.h
@@ -19,12 +19,12 @@
 #include <string>
 #include <vector>
 
-#include "fcp/aggregation/protocol/federated_compute_checkpoint_builder.h"
 #include "fcp/client/engine/common.h"
 #include "fcp/client/engine/example_iterator_factory.h"
 #include "fcp/client/example_iterator_query_recorder.h"
 #include "fcp/client/flags.h"
 #include "fcp/client/opstats/opstats_logger.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/protocol/federated_compute_checkpoint_builder.h"
 
 namespace fcp {
 namespace client {
@@ -48,7 +48,7 @@ class ExampleQueryPlanEngine {
  private:
   std::vector<ExampleIteratorFactory*> example_iterator_factories_;
   ::fcp::client::opstats::OpStatsLogger* opstats_logger_;
-  ::fcp::aggregation::FederatedComputeCheckpointBuilderFactory
+  tensorflow_federated::aggregation::FederatedComputeCheckpointBuilderFactory
       federated_compute_checkpoint_builder_factory_;
   const Flags& flags_;
   ExampleIteratorQueryRecorder* example_iterator_query_recorder_;
diff --git a/fcp/client/engine/example_query_plan_engine_test.cc b/fcp/client/engine/example_query_plan_engine_test.cc
index bc3c87a..a6ec9a9 100644
--- a/fcp/client/engine/example_query_plan_engine_test.cc
+++ b/fcp/client/engine/example_query_plan_engine_test.cc
@@ -33,11 +33,6 @@
 #include "absl/strings/cord.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
-#include "fcp/aggregation/core/tensor.h"
-#include "fcp/aggregation/core/tensor.pb.h"
-#include "fcp/aggregation/core/tensor_shape.h"
-#include "fcp/aggregation/protocol/checkpoint_header.h"
-#include "fcp/aggregation/testing/test_data.h"
 #include "fcp/client/client_runner.h"
 #include "fcp/client/engine/common.h"
 #include "fcp/client/engine/example_iterator_factory.h"
@@ -55,6 +50,11 @@
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/tstring.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor.pb.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/core/tensor_shape.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/protocol/checkpoint_header.h"
+#include "tensorflow_federated/cc/core/impl/aggregation/testing/test_data.h"
 
 namespace fcp {
 namespace client {
@@ -63,14 +63,14 @@ namespace {
 
 namespace tf = ::tensorflow;
 
-using ::fcp::aggregation::Tensor;
-using ::fcp::aggregation::TensorShape;
 using ::fcp::client::ExampleQueryResult;
 using ::google::internal::federated::plan::AggregationConfig;
 using ::google::internal::federated::plan::ClientOnlyPlan;
 using ::google::internal::federated::plan::Dataset;
 using ::google::internal::federated::plan::ExampleQuerySpec;
 using ::google::internal::federated::plan::ExampleSelector;
+using tensorflow_federated::aggregation::Tensor;
+using tensorflow_federated::aggregation::TensorShape;
 using ::testing::StrictMock;
 
 const char* const kCollectionUri = "app:/test_collection";
@@ -596,15 +596,18 @@ TEST_F(ExampleQueryPlanEngineTest,
 
   std::string header;
   ASSERT_TRUE(stream.ReadString(&header, 4));
-  ASSERT_EQ(header, aggregation::kFederatedComputeCheckpointHeader);
+  ASSERT_EQ(
+      header,
+      tensorflow_federated::aggregation::kFederatedComputeCheckpointHeader);
 
-  absl::StatusOr<Tensor> int_tensor =
-      Tensor::Create(aggregation::DT_INT64, TensorShape({2}),
-                     aggregation::CreateTestData<uint64_t>({42, 24}));
+  absl::StatusOr<Tensor> int_tensor = Tensor::Create(
+      tensorflow_federated::aggregation::DT_INT64, TensorShape({2}),
+      tensorflow_federated::aggregation::CreateTestData<uint64_t>({42, 24}));
   ASSERT_OK(int_tensor.status());
   absl::StatusOr<Tensor> string_tensor = Tensor::Create(
-      aggregation::DT_STRING, TensorShape({2}),
-      aggregation::CreateTestData<absl::string_view>({"value1", "value2"}));
+      tensorflow_federated::aggregation::DT_STRING, TensorShape({2}),
+      tensorflow_federated::aggregation::CreateTestData<absl::string_view>(
+          {"value1", "value2"}));
   ASSERT_OK(string_tensor.status());
 
   uint32_t name_size1;
@@ -612,7 +615,7 @@ TEST_F(ExampleQueryPlanEngineTest,
   std::string name1;
   ASSERT_TRUE(stream.ReadString(&name1, name_size1));
 
-  aggregation::Tensor& t1 =
+  tensorflow_federated::aggregation::Tensor& t1 =
       name1 == kOutputIntTensorName ? *int_tensor : *string_tensor;
   uint32_t tensor_size1;
   ASSERT_TRUE(stream.ReadVarint32(&tensor_size1));
@@ -625,7 +628,7 @@ TEST_F(ExampleQueryPlanEngineTest,
   std::string name2;
   ASSERT_TRUE(stream.ReadString(&name2, name_size2));
 
-  aggregation::Tensor& t2 =
+  tensorflow_federated::aggregation::Tensor& t2 =
       name2 == kOutputIntTensorName ? *int_tensor : *string_tensor;
   uint32_t tensor_size2;
   ASSERT_TRUE(stream.ReadVarint32(&tensor_size2));
diff --git a/fcp/demo/README.md b/fcp/demo/README.md
index 49cc687..302a6ed 100644
--- a/fcp/demo/README.md
+++ b/fcp/demo/README.md
@@ -213,7 +213,8 @@ Contains handlers for the Federated Compute protocol's
 [Aggregations](../protos/federatedcompute/aggregations.proto)
 service. This service is responsible for aggregating client-reported data using
 the
-[simple Aggregation Protocol](../aggregation/protocol/simple_aggregation)
+
+[simple Aggregation Protocol](https://github.com/tensorflow/federated/tree/main/tensorflow_federated/cc/core/impl/aggregation/protocol/simple_aggregation)
 library.
 
 Note that the demo platform does not currently contain an implementation of the