Add ToString method for both PjrtData and PjrtShardedData (#5265)

* Add ToString method for both PjrtData and PjrtShardedData * on cpu same config will become replicated, dont't check actual op sharding type

Author: JackCaoG

test/spmd/test_xla_sharding.py

@@ -37,6 +37,20 @@ def test_xla_sharded_tensor(self):
  # TODO(244003536) add more tests for XLAShardedTensror.
  self.assertTrue(isinstance(xst1, XLAShardedTensor))
 
+ def test_sharded_tensor_debug_info(self):
+ partition_spec = (0, 1)
+ xt1 = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]],
+ dtype=torch.float,
+ device=xm.xla_device())
+ xst1 = xs.mark_sharding(xt1, self._get_mesh((1, self.n_devices)),
+ partition_spec)
+
+ debug_info = torch_xla._XLAC._get_xla_tensor_debug_info(xst1.global_tensor)
+ self.assertIn('XLAShardedData', debug_info)
+ self.assertIn('Data Device: SPMD:0', debug_info)
+ self.assertIn('OpSharding:', debug_info)
+ self.assertIn('NumShards: %s' % (self.n_devices), debug_info)
+
  def test_xla_shards(self):
  num_element = self.n_devices
  mesh = self._get_mesh((self.n_devices,))

torch_xla/csrc/init_python_bindings.cpp

@@ -378,13 +378,11 @@ std::string GetXLATensorDebugInfo(const at::Tensor& tensor) {
  }
 
  torch::lazy::BackendDataPtr handle = xtensor->CurrentDataHandle();
- ss << "XLAData: ";
  if (handle) {
  auto data = UnwrapXlaData(handle);
- ss << "\n Data Device: " << data->device() << "\n";
- ss << " Data Shape: " << data->shape().ToString() << "\n";
+ ss << data->ToString();
  } else {
- ss << "None\n";
+ ss << "XLAData: None\n";
  }
 
  auto at_tensor = xtensor->CurrentTensorData();

torch_xla/csrc/runtime/computation_client.h

@@ -51,6 +51,8 @@ class ComputationClient {
 
  virtual bool HasValue() const = 0;
 
+ virtual std::string ToString() const = 0;
+
  private:
  std::string device_;
  xla::Shape shape_;

torch_xla/csrc/runtime/pjrt_computation_client.h

@@ -173,6 +173,20 @@ class PjRtComputationClient : public ComputationClient {
  return buffer != nullptr && !buffer->IsDeleted();
  };
 
+ std::string ToString() const override {
+ std::stringstream ss;
+ ss << "XLAData: \n";
+ ss << " Data Device: " << device() << "\n";
+ ss << " Data Shape: " << shape().ToString() << "\n";
+ ss << " Data Handle: ";
+ if (HasValue()) {
+ ss << reinterpret_cast<std::uintptr_t>(buffer.get()) << "\n";
+ } else {
+ ss << "None\n";
+ }
+ return ss.str();
+ }
+
  std::shared_ptr<xla::PjRtBuffer> buffer;
  };
 
@@ -212,6 +226,16 @@ class PjRtComputationClient : public ComputationClient {
  return true;
  }
 
+ std::string ToString() const override {
+ std::stringstream ss;
+ ss << "XLAShardedData: \n";
+ ss << " Data Device: " << device() << "\n";
+ ss << " Data Shape: " << shape().ToString() << "\n";
+ ss << " OpSharding: " << sharding.type() << "\n";
+ ss << " NumShards: " << shards.size() << "\n";
+ return ss.str();
+ }
+
  xla::OpSharding GetSharding() { return sharding; }
 
  std::vector<std::shared_ptr<PjRtData>> shards;