Minor build fixes and test skeleton

Author: darisoy

test/pjrt/test_experimental_pjrt_multi_cpu.py

@@ -0,0 +1,31 @@
+import os
+
+import torch
+import torch_xla
+from absl.testing import absltest, parameterized
+import torch_xla.core.xla_env_vars as xenv
+
+class TestExperimentalPjrtMultiCpu(parameterized.TestCase):
+    
+    def setUp(self):
+        pjrt.set_device_type('CPU')
+
+        os.environ.pop(xenv.CPU_ASYNC_CLIENT, None)
+        os.environ.pop(xenv.CPU_NUM_DEVICES, None)
+
+    def test_default_cpu_device(self):
+        devices_per_process = pjrt.run_multiprocess(xm.xla_device)
+        print(devices_per_process)
+
+    def test_multi_cpu_devices(self):
+        os.environ.update({
+            xenv.CPU_ASYNC_CLIENT: True,
+            xenv.CPU_NUM_DEVICES: 4,
+        })
+        devices_per_process = pjrt.run_multiprocess(xm.xla_device)
+        print(devices_per_process)
+
+
+
+if __name__ == '__main__':
+  absltest.main()

third_party/xla_client/BUILD

@@ -150,6 +150,7 @@ cc_library(
         "//tensorflow/compiler/xla/pjrt:cpu_device",
         "//tensorflow/compiler/xla/pjrt:tpu_client",
         "//tensorflow/compiler/xla/pjrt:pjrt_client",
+        "//tensorflow/compiler/xla/pjrt:tfrt_cpu_pjrt_client",
         "//tensorflow/compiler/xla/rpc:grpc_stub",
         "//tensorflow/compiler/xla/service:cpu_plugin",
         "//tensorflow/compiler/xla/service:platform_util",

third_party/xla_client/pjrt_computation_client.cc

@@ -10,6 +10,7 @@
 #include "tensorflow/compiler/xla/pjrt/cpu_device.h"
 #include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
 #include "tensorflow/compiler/xla/pjrt/tpu_client.h"
+#include "tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.h"
 #include "tensorflow/compiler/xla/shape.h"
 #include "tensorflow/compiler/xla/xla_client/computation_client.h"
 #include "tensorflow/compiler/xla/xla_client/debug_macros.h"