source ci/official/envs/linux_x86
TFCI_BAZEL_COMMON_ARGS="--repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config release_cpu_linux --config=tpu"
TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu
TFCI_BUILD_PIP_PACKAGE_ARGS="--repo_env=WHEEL_NAME=tensorflow_tpu"
TFCI_LIB_SUFFIX="-tpu-linux-x86_64"
TFCI_WHL_BAZEL_TEST_ENABLE=0
TFCI_WHL_IMPORT_TEST_ENABLE=0
TFCI_WHL_SIZE_LIMIT=580M

# bazel test invocation as normal.
test:nonpip_filters --test_tag_filters=-no_oss,-oss_serial,-gpu,-tpu,-benchmark-test,-v1only,-no_aarch64,-no_oss_py38,-no_oss_py39,-no_oss_py310
test:nonpip_filters --build_tag_filters=-no_oss,-oss_serial,-gpu,-tpu,-benchmark-test,-v1only,-no_aarch64,-no_oss_py38,-no_oss_py39,-no_oss_py310

# bazel test invocation as normal.
test:nonpip_filters --test_tag_filters=-no_oss,-oss_serial,-gpu,-tpu,-benchmark-test,-v1only,-no_aarch64,-no_oss_py38,-no_oss_py39,-no_oss_py310
test:nonpip_filters --build_tag_filters=-no_oss,-oss_serial,-gpu,-tpu,-benchmark-test,-v1only,-no_aarch64,-no_oss_py38,-no_oss_py39,-no_oss_py310

# AWS SDK must be compiled in release mode. see: https://github.com/tensorflow/tensorflow/issues/37498
build:dbg --copt -DDEBUG_BUILD

# Config to build TF TPU
build:tpu --define=with_tpu_support=true
build:tpu --define=framework_shared_object=true
build:tpu --copt=-DLIBTPU_ON_GCE
build:tpu --define=enable_mlir_bridge=true

build:tensorrt --repo_env TF_NEED_TENSORRT=1

build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain

    if (device_type == "TPU") {
      has_tpu = true;
      break;
    }
  }
  if (has_tpu) {
    BasicTestsForTwoDevices(context.get(),
                            "/job:localhost/replica:0/task:0/device:TPU:0",
                            "/job:localhost/replica:0/task:0/device:TPU:0");
  }
}

TEST(PARALLEL_DEVICE, TestExplicitCopies) {

)

# This flag enables experimental TPU support
bool_flag(
    name = "enable_tpu_support",
    build_setting_default = False,
)

config_setting(
    name = "with_tpu_support_define",
    define_values = {"with_tpu_support": "true"},
    visibility = ["//visibility:public"],
)

config_setting(
    name = "with_tpu_support_flag",
    flag_values = {":enable_tpu_support": "True"},

         * It has an added advantage since you can you easily switch to different hardware accelerators (cpu, gpu, tpu) as per the task.
         * All you need is a good internet connection and you are all set.
      * Try to build TF from sources by changing CPU optimization flags.
   
   *Please let us know if this helps.*
   

              TFE_TensorHandle* axis);

// If there is a device of type `device_type`, returns true
// and sets 'device_name' accordingly.
// `device_type` must be either "GPU" or "TPU".
bool GetDeviceName(TFE_Context* ctx, tensorflow::string* device_name,
                   const char* device_type);

// Create a ServerDef with the given `job_name` and add `num_tasks` tasks in it.

    const TFE_OpAttrs* attributes, int expected_max_outputs,
    TF_Status* status) {
  absl::optional<std::vector<MaybeParallelTensorOwned>> result;
  // TODO(allenl): We should remove "TPU" from these op names at the very least,
  // or consider other ways of packing/unpacking parallel tensors.
  if (operation_name == std::string("TPUReplicatedInput")) {

*   `tf.tpu.experimental.embedding`:

    *   `tf.tpu.experimental.embedding.FeatureConfig` now takes an additional
        argument `output_shape` which can specify the shape of the output
        activation for the feature.
    *   `tf.tpu.experimental.embedding.TPUEmbedding` now has the same behavior
        as `tf.tpu.experimental.embedding.serving_embedding_lookup` which can