Clean up Tensor and Shape stringification (flashlight#808)

Summary:
Pull Request resolved: flashlight#808

See title. Add `Tensor::toString()` and `Shape::toString()` explicitly and use in `operator<<` for `ArrayFireTensor`.

Reviewed By: benoitsteiner

Differential Revision: D33159894

fbshipit-source-id: bc6a32658191933b37662c44ac13561e1bd3e305

flashlight/fl/tensor/Shape.cpp

@@ -73,15 +73,6 @@ bool Shape::operator!=(const std::initializer_list<Dim>& other) const {
   return !(this->operator==(other));
 }
 
-std::ostream& operator<<(std::ostream& ostr, const Shape& s) {
-  ostr << "(";
-  for (size_t i = 0; i < s.ndim(); ++i) {
-    ostr << s.dim(i) << (i == s.ndim() - 1 ? "" : ", ");
-  }
-  ostr << ")";
-  return ostr;
-}
-
 const std::vector<Dim>& Shape::get() const {
   return dims_;
 }
@@ -90,4 +81,19 @@ std::vector<Dim>& Shape::get() {
   return dims_;
 };
 
+std::string Shape::toString() const {
+  std::stringstream ss;
+  ss << "(";
+  for (size_t i = 0; i < ndim(); ++i) {
+    ss << dim(i) << (i == ndim() - 1 ? "" : ", ");
+  }
+  ss << ")";
+  return ss.str();
+}
+
+std::ostream& operator<<(std::ostream& ostr, const Shape& s) {
+  ostr << s.toString();
+  return ostr;
+}
+
 } // namespace fl

flashlight/fl/tensor/Shape.h

@@ -113,6 +113,11 @@ class Shape {
    */
   const std::vector<Dim>& get() const;
   std::vector<Dim>& get();
+
+  /**
+   * Returns a string representation of the Shape
+   */
+  std::string toString() const;
 };
 
 /**

flashlight/fl/tensor/TensorAdapter.h

@@ -212,6 +212,12 @@ class TensorAdapterBase {
    */
   virtual void* getContext() = 0;
 
+  /**
+   * Return a string representation of a Tensor. Not intended to be portable
+   * across backends.
+   */
+  virtual std::string toString() = 0;
+
   /**
    * Write a string representation of a tensor to an output stream.
    */

flashlight/fl/tensor/TensorBase.cpp

@@ -240,6 +240,10 @@ void* Tensor::getContext() const {
   return impl_->getContext();
 }
 
+std::string Tensor::toString() const {
+  return impl_->toString();
+}
+
 std::ostream& Tensor::operator<<(std::ostream& ostr) const {
   return impl_->operator<<(ostr);
 }

flashlight/fl/tensor/TensorBase.h

@@ -521,6 +521,15 @@ class Tensor {
    */
   void* getContext() const;
 
+  /**
+   * Returns a string representation of a Tensor. NOTE: This is
+   * backend-dependent. See Flashlight's serialization utilities for ways to
+   * serialize Tensors that are portable across Tensor backends.
+   *
+   * @return a string representation of the Tensor.
+   */
+  std::string toString() const;
+
   /**
    * Write a string representation of a tensor to an output stream.
    */
@@ -1391,7 +1400,7 @@ Tensor all(
 /**
  * Write a string representation of a tensor to an output stream.
  */
-std::ostream& operator<<(std::ostream& ostr, const Tensor& s);
+std::ostream& operator<<(std::ostream& ostr, const Tensor& t);
 
 /**
  * Print a string representation of a tensor to standard out.

flashlight/fl/tensor/backend/af/ArrayFireTensor.cpp

@@ -336,8 +336,12 @@ void* ArrayFireTensor::getContext() {
   return nullptr; // noop
 }
 
+std::string ArrayFireTensor::toString() {
+  return std::string(af::toString("ArrayFireTensor", getHandle()));
+}
+
 std::ostream& ArrayFireTensor::operator<<(std::ostream& ostr) {
-  ostr << "ArrayFireTensor " << std::string(af::toString("", getHandle()));
+  ostr << this->toString();
   return ostr;
 }
 

flashlight/fl/tensor/backend/af/ArrayFireTensor.h

@@ -200,6 +200,7 @@ class ArrayFireTensor : public TensorAdapterBase {
   Tensor asContiguousTensor() override;
   void setContext(void* context) override; // noop
   void* getContext() override; // noop
+  std::string toString() override;
   std::ostream& operator<<(std::ostream& ostr) override;
 
   /******************** Assignment Operators ********************/

flashlight/fl/test/tensor/ShapeTest.cpp

@@ -71,3 +71,18 @@ TEST(ShapeTest, Indexing) {
   ASSERT_EQ(a[3], 2);
   ASSERT_THROW(a[4], std::invalid_argument);
 }
+
+TEST(ShapeTest, string) {
+  auto checkShapeStrEqual = [](const Shape& s, const std::string& str) -> void {
+    auto sStr = s.toString();
+    ASSERT_EQ(sStr, str);
+    std::stringstream ss;
+    ss << sStr;
+    ASSERT_EQ(sStr, ss.str());
+  };
+
+  checkShapeStrEqual(Shape({3, 4, 7, 9}), "(3, 4, 7, 9)");
+  checkShapeStrEqual(Shape({}), "()");
+  checkShapeStrEqual(Shape({0}), "(0)");
+  checkShapeStrEqual(Shape({7, 7, 7, 7, 7, 7, 7}), "(7, 7, 7, 7, 7, 7, 7)");
+}

flashlight/fl/test/tensor/TensorBaseTest.cpp

@@ -67,11 +67,13 @@ TEST(TensorBaseTest, fromScalar) {
   ASSERT_EQ(a.shape(), Shape({}));
 }
 
-TEST(TensorBaseTest, ostream) {
+TEST(TensorBaseTest, string) {
   // Different backends might print tensors differently - check for consistency
   // across two identical tensors
   auto a = fl::full({3, 4, 5}, 6.);
   auto b = fl::full({3, 4, 5}, 6.);
+  ASSERT_EQ(a.toString(), b.toString());
+
   std::stringstream ssa, ssb;
   ssa << a;
   ssb << b;