# Copyright 2026 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.
# ==============================================================================

load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = [
        # copybara:uncomment "//learning/brain/tfrt:__subpackages__",
        # copybara:uncomment "//third_party/neptune:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/tfrt:__subpackages__",
        "//tensorflow/core/tfrt/mlrt:__subpackages__",
    ],
)

tf_proto_library(
    name = "ifrt_config_proto",
    srcs = ["ifrt_config.proto"],
    protodeps = [
        "@xla//xla:xla_data_proto",
        "//tensorflow/core/framework:tensor_proto",
    ],
    visibility = ["//visibility:public"],
)

cc_library(
    name = "ifrt_serving_core_selector",
    srcs = ["ifrt_serving_core_selector.cc"],
    hdrs = ["ifrt_serving_core_selector.h"],
    deps = [
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@xla//xla/tsl/framework:serving_device_selector",
    ],
)

tf_cc_test(
    name = "ifrt_serving_core_selector_test",
    srcs = ["ifrt_serving_core_selector_test.cc"],
    tags = ["no_oss"],
    deps = [
        ":ifrt_serving_core_selector",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
        "@xla//xla/tsl/framework:serving_device_selector",
        "@xla//xla/tsl/framework/test_util:mock_serving_device_selector",
    ],
)

cc_library(
    name = "grid",
    srcs = ["grid.cc"],
    hdrs = ["grid.h"],
    deps = [
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "ifrt_device_utils",
    srcs = ["ifrt_device_utils.cc"],
    hdrs = ["ifrt_device_utils.h"],
    deps = [
        ":grid",
        ":ifrt_config_proto_cc",
        "//tensorflow/compiler/tf2xla:host_compute_metadata_proto_cc",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/core:core_cpu_base",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:attribute_map",
        "@xla//xla/service:device_assignment",
        "@xla//xla/tsl/platform:statusor",
    ],
)

# copybara:uncomment_begin(PjRt-IFRT TPU client is not available in OSS)
# # TODO(siqiaowu): ifrt_device_utils_test couples with TPU assumptions ("coords"
# # and "core_on_chip" device attributes) and requires a TPU test client, which is
# # not available in OSS.
# tf_cc_test(
#     name = "ifrt_device_utils_test",
#     srcs = [
#         "ifrt_device_utils_test.cc",
#     ],
#     tags = [
#         "no_oss",
#         "requires-tpu",
#     ],
#     deps = [
#         ":ifrt_device_utils",
#         "//learning/brain/research/ifrt/pjrt_ifrt/pjrt_tpu:pjrt_tpu_client_test_lib",
#         "@com_google_googletest//:gtest_main",
#         "@com_google_absl//absl/status",
#         "@com_google_absl//absl/status:status_matchers",
#         "@xla//xla/python/ifrt",
#         "@xla//xla/python/ifrt:attribute_map",
#         "@xla//xla/python/ifrt:test_util",
#         "@xla//xla/service:device_assignment",
#         "@xla//xla/tsl/platform:status_matchers",
#         "@xla//xla/tsl/platform:statusor",
#     ],
# )
# copybara:uncomment_end

cc_library(
    name = "ifrt_serving_executable",
    srcs = ["ifrt_serving_executable.cc"],
    hdrs = ["ifrt_serving_executable.h"],
    deps = [
        ":ifrt_config_proto_cc",
        ":ifrt_device_utils",
        ":ifrt_loaded_variable_registry",
        ":ifrt_loaded_variable_utils",
        ":ifrt_persistent_compilation_cache",
        ":ifrt_restore_tensor_registry",
        ":ifrt_serving_core_selector",
        ":ifrt_tensor_utils",
        ":sharding_utils",
        ":tf_host_callback",
        "//tensorflow/compiler/mlir/tensorflow",
        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
        "//tensorflow/compiler/mlir/tfrt:export",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:extract_callback",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:host_compute_metadata_proto_cc",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/core:core_cpu_base",
        "//tensorflow/core:framework",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/hash",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/log:vlog_is_on",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:Support",
        "@tf_runtime//:hostcontext",
        "@tsl//tsl/platform:tstring",
        "@tsl//tsl/profiler/lib:traceme",
        "@xla//xla:shape_util",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/hlo/translate:stablehlo",
        "@xla//xla/pjrt:host_callback",
        "@xla//xla/pjrt:pjrt_client",
        "@xla//xla/pjrt:pjrt_compiler",
        "@xla//xla/pjrt:pjrt_executable",
        "@xla//xla/pjrt:pjrt_layout",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt/hlo:hlo_program",
        "@xla//xla/python/pjrt_ifrt",
        "@xla//xla/python/pjrt_ifrt:pjrt_layout",
        "@xla//xla/service:device_assignment",
        "@xla//xla/service:dump",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/concurrency:ref_count",
        "@xla//xla/tsl/framework:serving_device_selector",
        "@xla//xla/tsl/lib/monitoring:counter",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:errors",
        "@xla//xla/tsl/platform:statusor",
    ],
)

cc_library(
    name = "ifrt_executable_registry",
    srcs = ["ifrt_executable_registry.cc"],
    hdrs = ["ifrt_executable_registry.h"],
    deps = [
        ":ifrt_serving_executable",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
    ],
)

cc_library(
    name = "ifrt_restore_tensor_registry",
    srcs = ["ifrt_restore_tensor_registry.cc"],
    hdrs = ["ifrt_restore_tensor_registry.h"],
    deps = [
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/core:framework",
        "//tensorflow/core/framework:tensor",
        "//tensorflow/core/framework:types_proto_cc",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@xla//xla/tsl/concurrency:future",
    ],
)

cc_library(
    name = "ifrt_persistent_compilation_cache",
    srcs = ["ifrt_persistent_compilation_cache.cc"],
    hdrs = ["ifrt_persistent_compilation_cache.h"],
    deps = [
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
        "@xla//xla/pjrt:pjrt_executable",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt/hlo:hlo_program",
        "@xla//xla/python/pjrt_ifrt:xla_ifrt",
        "@xla//xla/tsl/concurrency:ref_count",
    ],
)

cc_library(
    name = "ifrt_loaded_variable_registry",
    srcs = ["ifrt_loaded_variable_registry.cc"],
    hdrs = ["ifrt_loaded_variable_registry.h"],
    deps = [
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/hash",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@xla//xla:shape_util",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/python/ifrt",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/platform:statusor",
    ],
)

tf_cc_test(
    name = "ifrt_loaded_variable_registry_test",
    srcs = ["ifrt_loaded_variable_registry_test.cc"],
    tags = ["no_oss"],
    deps = [
        ":ifrt_loaded_variable_registry",
        "@com_google_absl//absl/hash",
        "@com_google_absl//absl/hash:hash_testing",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
        "@xla//xla:shape_util",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/python/ifrt",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/lib/core:status_test_util",
        "@xla//xla/tsl/platform:status_matchers",
    ],
)

cc_library(
    name = "ifrt_model_restore_context",
    hdrs = ["ifrt_model_restore_context.h"],
    visibility = [
        "//learning/brain/tfrt:__subpackages__",
        "//learning/serving/servables/tfrt:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/tfrt:__subpackages__",
    ],
    deps = [
        ":checkpoint_loader",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "ifrt_model_context",
    srcs = ["ifrt_model_context.cc"],
    hdrs = ["ifrt_model_context.h"],
    visibility = [
        "//learning/brain/tfrt:__subpackages__",
        "//learning/serving/servables/tfrt:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/tfrt:__subpackages__",
        "//tensorflow/core/tfrt/mlrt:__subpackages__",
    ],
    deps = [
        ":ifrt_config_proto_cc",
        ":ifrt_executable_registry",
        ":ifrt_loaded_variable_registry",
        ":ifrt_persistent_compilation_cache",
        ":ifrt_restore_tensor_registry",
        ":ifrt_serving_core_selector",
        ":sharding_utils",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/core:core_cpu_base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tf_runtime//:hostcontext",
        "@tsl//tsl/platform:protobuf",
        "@xla//xla/pjrt:pjrt_executable",
        "@xla//xla/python/ifrt",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:errors",
    ],
)

cc_library(
    name = "ifrt_tensor_utils",
    srcs = ["ifrt_tensor_utils.cc"],
    hdrs = ["ifrt_tensor_utils.h"],
    deps = [
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/core:framework",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/pjrt_ifrt:pjrt_dtype",
        "@xla//xla/tsl/platform:errors",
        "@xla//xla/tsl/platform:statusor",
    ],
)

cc_library(
    name = "sharding_utils",
    srcs = [
        "sharding_utils.cc",
    ],
    hdrs = [
        "sharding_utils.h",
    ],
    deps = [
        ":ifrt_tensor_utils",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/core:framework",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core/tpu/kernels:sharding_utils",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:btree",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
        "@xla//xla:shape_util",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/pjrt:pjrt_client",
        "@xla//xla/pjrt:pjrt_layout",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/pjrt_ifrt",
        "@xla//xla/python/pjrt_ifrt:xla_ifrt",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/concurrency:ref_count",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:errors",
        "@xla//xla/tsl/platform:statusor",
    ],
)

cc_library(
    name = "ifrt_loaded_variable_utils",
    srcs = ["ifrt_loaded_variable_utils.cc"],
    hdrs = ["ifrt_loaded_variable_utils.h"],
    deps = [
        ":ifrt_config_proto_cc",
        ":ifrt_loaded_variable_registry",
        ":ifrt_restore_tensor_registry",
        ":sharding_utils",
        ":undonatable_buffer_converter",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/core:framework",
        "//tensorflow/core/platform:context",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@tf_runtime//:hostcontext",
        "@xla//xla:future",
        "@xla//xla:shape_util",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/python/ifrt",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:errors",
        "@xla//xla/tsl/platform:statusor",
    ],
)

cc_library(
    name = "undonatable_buffer_converter",
    srcs = ["undonatable_buffer_converter.cc"],
    hdrs = ["undonatable_buffer_converter.h"],
    deps = [
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@xla//xla:shape_util",
        "@xla//xla/pjrt:common_pjrt_client",
        "@xla//xla/pjrt:device_event",
        "@xla//xla/pjrt:pjrt_client",
        "@xla//xla/pjrt:raw_buffer",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/pjrt_ifrt",
        "@xla//xla/tsl/platform:errors",
        "@xla//xla/tsl/platform:statusor",
    ],
)

cc_library(
    name = "tf_host_callback",
    srcs = ["tf_host_callback.cc"],
    hdrs = ["tf_host_callback.h"],
    deps = [
        "//tensorflow/c/eager:abstract_tensor_handle",
        "//tensorflow/c/eager:immediate_execution_context",
        "//tensorflow/c/eager:immediate_execution_operation",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/core:core_cpu_base",
        "//tensorflow/core:framework",
        "//tensorflow/core:framework_internal",
        "//tensorflow/core/common_runtime/eager:context",
        "//tensorflow/core/common_runtime/eager:core_no_xla",
        "//tensorflow/core/common_runtime/eager:tensor_handle",
        "//tensorflow/core/framework:function_proto_cc",
        "//tensorflow/core/protobuf:for_core_protos_cc",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/cleanup",
        "@com_google_absl//absl/container:fixed_array",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:casts",
        "@tsl//tsl/platform:refcount",
        "@tsl//tsl/profiler/lib:traceme",
        "@xla//xla/tsl/platform:errors",
    ],
)

cc_library(
    name = "ifrt_serving_executable_test_util",
    testonly = True,
    srcs = ["ifrt_serving_executable_test_util.cc"],
    hdrs = ["ifrt_serving_executable_test_util.h"],
    data = [
        "//tensorflow/core/tfrt/ifrt/testdata",
    ],
    deps = [
        ":ifrt_loaded_variable_registry",
        ":ifrt_persistent_compilation_cache",
        ":ifrt_restore_tensor_registry",
        ":ifrt_serving_core_selector",
        ":ifrt_serving_executable",
        ":sharding_utils",
        ":tf_host_callback",
        "//tensorflow/compiler/mlir/tensorflow",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/core:core_cpu_base",
        "//tensorflow/core/framework:tensor",
        "//tensorflow/core/framework:types_proto_cc",
        "//tensorflow/core/platform:resource_loader",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@llvm-project//mlir:AllPassesAndDialects",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:RegisterAllDialects",
        "@tf_runtime//:hostcontext",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:test_util",
        "@xla//xla/tsl/framework/test_util:mock_serving_device_selector",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:status",
    ],
)

tf_cc_test(
    name = "ifrt_restore_tensor_registry_test",
    srcs = ["ifrt_restore_tensor_registry_test.cc"],
    tags = ["no_oss"],
    deps = [
        ":ifrt_restore_tensor_registry",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/core:framework",
        "//tensorflow/core/framework:tensor",
        "//tensorflow/core/framework:tensor_testutil",
        "//tensorflow/core/framework:types_proto_cc",
        "@com_google_absl//absl/base:log_severity",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/lib/core:status_test_util",
        "@xla//xla/tsl/platform:status_matchers",
        "@xla//xla/tsl/platform:statusor",
    ],
)

tf_cc_test(
    name = "ifrt_loaded_variable_utils_test",
    srcs = ["ifrt_loaded_variable_utils_test.cc"],
    tags = ["no_oss"],
    deps = [
        ":ifrt_config_proto_cc",
        ":ifrt_loaded_variable_registry",
        ":ifrt_loaded_variable_utils",
        ":ifrt_restore_tensor_registry",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/core:framework",
        "//tensorflow/core/framework:tensor_matcher",
        "//tensorflow/core/framework:tensor_testutil",
        "//tensorflow/core/framework:types_proto_cc",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
        "@tf_runtime//:hostcontext",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:test_util",
        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/concurrency:ref_count",
        "@xla//xla/tsl/lib/core:status_test_util",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:statusor",
        "@xla//xla/tsl/platform:test",
    ],
)

tf_cc_test(
    name = "undonatable_buffer_converter_test",
    srcs = ["undonatable_buffer_converter_test.cc"],
    deps = [
        ":sharding_utils",
        ":undonatable_buffer_converter",
        "//tensorflow/core:framework",
        "//tensorflow/core/framework:tensor_matcher",
        "//tensorflow/core/framework:tensor_testutil",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/pjrt:common_pjrt_client",
        "@xla//xla/pjrt:pjrt_client",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:test_util",
        "@xla//xla/python/pjrt_ifrt",
        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:statusor",
    ],
)

tf_cc_test(
    name = "tf_host_callback_test",
    srcs = ["tf_host_callback_test.cc"],
    tags = ["no_oss"],
    deps = [
        ":tf_host_callback",
        "//tensorflow/cc:const_op",
        "//tensorflow/cc:function_ops",
        "//tensorflow/cc:functional_ops",
        "//tensorflow/cc:math_ops",
        "//tensorflow/cc:ops",
        "//tensorflow/cc:resource_variable_ops",
        "//tensorflow/cc:scope",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/core:core_cpu_base",
        "//tensorflow/core:framework",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core:tensorflow",
        "//tensorflow/core/framework:tensor_matcher",
        "//tensorflow/core/framework:tensor_testutil",
        "//tensorflow/core/kernels:constant_op",
        "//tensorflow/core/kernels:math",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/tsl/platform:errors",
        "@xla//xla/tsl/platform:statusor",
    ],
)

tf_cc_test(
    name = "sharding_utils_test",
    srcs = ["sharding_utils_test.cc"],
    tags = ["no_oss"],
    deps = [
        ":ifrt_tensor_utils",
        ":pjrt_cpu_client_test_lib",
        ":sharding_utils",
        "//tensorflow/core:framework",
        "//tensorflow/core:test",
        "//tensorflow/core/framework:tensor_matcher",
        "//tensorflow/core/framework:tensor_testutil",
        "//tensorflow/core/framework:types_proto_cc",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:ml_dtypes",
        "@xla//xla:shape_util",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/hlo/ir:hlo",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:test_util",
        "@xla//xla/python/pjrt_ifrt:xla_ifrt",
        "@xla//xla/tsl/concurrency:ref_count",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:statusor",
    ],
)

tf_cc_test(
    name = "ifrt_serving_executable_test",
    srcs = [
        "ifrt_serving_executable_test.cc",
    ],
    data = [
        "//tensorflow/core/tfrt/ifrt/testdata",
    ],
    tags = ["no_oss"],
    deps = [
        ":ifrt_restore_tensor_registry",
        ":ifrt_serving_executable",
        ":ifrt_serving_executable_test_util",
        ":sharding_utils",
        "//tensorflow/compiler/mlir/tensorflow",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/core:core_cpu_base",
        "//tensorflow/core:framework",
        "//tensorflow/core:portable_tensorflow_test_lib",
        "//tensorflow/core:test",
        "//tensorflow/core/common_runtime/eager:core",
        "//tensorflow/core/framework:tensor",
        "//tensorflow/core/framework:tensor_matcher",
        "//tensorflow/core/framework:tensor_testutil",
        "//tensorflow/core/framework:types_proto_cc",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/log:globals",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:AllPassesAndDialects",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:Support",
        "@tf_runtime//:basic_kernels_alwayslink",
        "@tf_runtime//:core_runtime_alwayslink",
        "@tf_runtime//:test_kernels_alwayslink",
        "@tf_runtime//backends/cpu:core_runtime_alwayslink",
        "@tf_runtime//backends/cpu:tf_ops_alwayslink",
        "@tsl//tsl/platform:tstring",
        "@xla//xla:shape_util",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/pjrt:mlir_to_hlo",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:test_util",
        "@xla//xla/python/ifrt/hlo:hlo_program",
        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
        "@xla//xla/python/pjrt_ifrt:xla_ifrt",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/framework:serving_device_selector",
        "@xla//xla/tsl/framework/test_util:mock_serving_device_selector",
        "@xla//xla/tsl/lib/monitoring:cell_reader",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:statusor",
    ],
)

tf_cc_test(
    name = "ifrt_executable_registry_test",
    srcs = [
        "ifrt_executable_registry_test.cc",
    ],
    data = [
        "//tensorflow/core/tfrt/ifrt/testdata",
    ],
    tags = ["no_oss"],
    deps = [
        ":ifrt_executable_registry",
        ":ifrt_loaded_variable_registry",
        ":ifrt_restore_tensor_registry",
        ":ifrt_serving_executable",
        ":sharding_utils",
        ":tf_host_callback",
        "//tensorflow/compiler/mlir/tensorflow",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/core:core_cpu_base",
        "//tensorflow/core:lib",
        "//tensorflow/core:portable_tensorflow_test_lib",
        "//tensorflow/core:test",
        "//tensorflow/core/common_runtime/eager:core",
        "//tensorflow/core/framework:tensor",
        "//tensorflow/core/framework:types_proto_cc",
        "//tensorflow/core/platform:resource_loader",
        "//tensorflow/core/platform:status_matchers",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:AllPassesAndDialects",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:RegisterAllDialects",
        "@tf_runtime//:basic_kernels_alwayslink",
        "@tf_runtime//:core_runtime_alwayslink",
        "@tf_runtime//:hostcontext",
        "@tf_runtime//:test_kernels_alwayslink",
        "@tf_runtime//backends/cpu:core_runtime_alwayslink",
        "@tf_runtime//backends/cpu:tf_ops_alwayslink",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:test_util",
        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
        "@xla//xla/tsl/platform:env",
        "@xla//xla/tsl/platform:statusor",
    ],
)

cc_library(
    name = "checkpoint_loader",
    srcs = ["checkpoint_loader.cc"],
    hdrs = ["checkpoint_loader.h"],
    visibility = [
        "//learning/brain/tfrt:__subpackages__",
        "//learning/serving/servables/tfrt:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/tfrt:__subpackages__",
    ],
    deps = [
        ":ifrt_loaded_variable_utils",
        ":ifrt_restore_tensor_registry",
        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
        "//tensorflow/core:framework",
        "//tensorflow/core:framework_internal",
        "//tensorflow/core:lib",
        "//tensorflow/core/common_runtime:function",
        "//tensorflow/core/framework:attr_value_proto_cc",
        "//tensorflow/core/framework:node_def_util",
        "//tensorflow/core/framework:tensor",
        "//tensorflow/core/framework:types_proto_cc",
        "//tensorflow/core/platform:refcount",
        "//tensorflow/core/protobuf:for_core_protos_cc",
        "//tensorflow/core/tfrt/fallback:fallback_state",
        "//tensorflow/core/tfrt/fallback:op_kernel_runner",
        "//tensorflow/core/tfrt/mlrt/kernel:context",
        "//tensorflow/core/tfrt/mlrt/kernel:kernel_runner_utils",
        "//tensorflow/core/tfrt/mlrt/kernel:shard_restore_util",
        "//tensorflow/core/tfrt/utils:fallback_tensor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
        "@tf_runtime//:hostcontext",
        "@tsl//tsl/platform:tstring",
        "@xla//xla/tsl/concurrency:future",
        "@xla//xla/tsl/platform:errors",
        "@xla//xla/tsl/platform:statusor",
    ],
)

cc_library(
    name = "pjrt_cpu_client_test_lib",
    testonly = True,
    srcs = ["pjrt_cpu_client_test_lib.cc"],
    deps = [
        "@com_google_absl//absl/status:statusor",
        "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options",
        "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
        "@xla//xla/python/ifrt",
        "@xla//xla/python/ifrt:test_util",
        "@xla//xla/python/pjrt_ifrt",
        "@xla//xla/tsl/platform:statusor",
    ],
    alwayslink = True,
)
