[TF-76][Python Interop] add PythonConvertible conformance to TensorShape (#23762)

* [TF-76] add TensorShape conformance to PythonConvertible

Author: realdoug

stdlib/public/TensorFlow/TensorShape.swift

@@ -10,6 +10,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Python
+
 // NOTE: it may be possible to edit `TensorShape` to support "labeled tensors".
 // Dimensions may be either an Int32 or an enum representing a label.
 
@@ -158,3 +160,20 @@ extension TensorShape : Codable {
     self.init(dimensions)
   }
 }
+
+extension TensorShape : PythonConvertible {
+  public var pythonObject: PythonObject {
+    return dimensions.pythonObject
+  }
+
+  public init?(_ pythonObject: PythonObject) {
+    let hasLen = Bool(Python.hasattr(pythonObject, "__len__"))
+    if(hasLen == true) {
+      guard let array = [Int32](pythonObject) else { return nil }
+      self.init(array)
+    } else {
+      guard let num = Int32(pythonObject) else { return nil }
+      self.init(num)
+    }
+  }
+}

test/Python/python_runtime.swift

@@ -66,7 +66,7 @@ PythonRuntimeTestSuite.testWithLeakChecking("PythonList") {
 }
 
 PythonRuntimeTestSuite.testWithLeakChecking("PythonDict") {
-  let dict: PythonObject = ["a": 1, 1: 0.5]
+  let dict: PythonObject = ["a" : 1, 1 : 0.5]
   expectEqual(2, Python.len(dict))
   expectEqual(1, dict["a"])
   expectEqual(0.5, dict[1])
@@ -210,7 +210,7 @@ PythonRuntimeTestSuite.testWithLeakChecking("Tuple") {
   let element1: PythonObject = 0
   let element2: PythonObject = "abc"
   let element3: PythonObject = [0, 0]
-  let element4: PythonObject = ["a": 0, "b": "c"]
+  let element4: PythonObject = ["a" : 0, "b" : "c"]
   let pair = PythonObject(tupleOf: element1, element2)
   let (pair1, pair2) = pair.tuple2
   expectEqual(element1, pair1)
@@ -258,6 +258,8 @@ PythonRuntimeTestSuite.testWithLeakChecking("ConvertibleFromPython") {
   let five: PythonObject = 5
   let half: PythonObject = 0.5
   let string: PythonObject = "abc"
+  let intArray: PythonObject = [2, 3]
+  let dict: PythonObject = ["abc" : 97]
 
   expectEqual(-1, Int(minusOne))
   expectEqual(-1, Int8(minusOne))
@@ -285,6 +287,9 @@ PythonRuntimeTestSuite.testWithLeakChecking("ConvertibleFromPython") {
 
   expectEqual("abc", String(string))
 
+  expectEqual([2, 3], Array(intArray))
+  expectEqual(["abc" : 97], Dictionary<String, Int32>(dict))
+
   expectNil(String(zero))
   expectNil(Int(string))
   expectNil(Double(string))
@@ -293,6 +298,8 @@ PythonRuntimeTestSuite.testWithLeakChecking("ConvertibleFromPython") {
 PythonRuntimeTestSuite.testWithLeakChecking("PythonConvertible") {
   let minusOne: PythonObject = -1
   let five: PythonObject = 5
+  let intArray: PythonObject = [2, 3]
+  let dict: PythonObject = ["abc" : 7]
 
   expectEqual(minusOne, Int(-1).pythonObject)
   expectEqual(minusOne, Int8(-1).pythonObject)
@@ -309,6 +316,9 @@ PythonRuntimeTestSuite.testWithLeakChecking("PythonConvertible") {
   expectEqual(five, UInt64(5).pythonObject)
   expectEqual(five, Float(5).pythonObject)
   expectEqual(five, Double(5).pythonObject)
+
+  expectEqual(intArray, [2, 3].pythonObject)
+  expectEqual(dict, ["abc" : 7].pythonObject)
 }
 
 PythonRuntimeTestSuite.testWithLeakChecking("Optional") {

test/TensorFlow/integration.swift

@@ -377,7 +377,7 @@ public func testResourceAndVariants() {
     // expected-error @+1 {{op named 'TensorDataSet' is not registered in TensorFlow}}
     #tfop("TensorDataSet", values,
           Toutput_types$dtype: [Float.tensorFlowDataType],
-          output_shapes: [TensorShape(1)])
+          output_shapes: [TensorShape([1])])
 
   // REGISTER_OP("Iterator")
   //     .Output("handle: resource")
@@ -388,7 +388,7 @@ public func testResourceAndVariants() {
   //     .SetShapeFn(shape_inference::ScalarShape);
   let iterator: ResourceHandle =
     #tfop("Iterator", shared_name: "foo", container: "bar",
-          output_types$dtype: [Float.tensorFlowDataType], output_shapes: [TensorShape(1)])
+          output_types$dtype: [Float.tensorFlowDataType], output_shapes: [TensorShape([1])])
 
   // REGISTER_OP("MakeIterator")
   //     .Input("dataset: variant")

test/TensorFlowRuntime/numpy_conversion.swift

@@ -52,6 +52,18 @@ NumpyConversionTests.test("shaped-array-conversion") {
                 array)
   }
 
+  let reshaped = np.reshape(numpyArrayInt32, [2, 3] as TensorShape)
+  if let array = expectNotNil(ShapedArray<Int32>(numpy: reshaped)) {
+    expectEqual(ShapedArray(shape: [2, 3], scalars: [1, 2, 3, 4, 5, 6]),
+                array)
+  }
+
+  let numpyArray1D = np.ones(28)
+  let reshaped3D = np.reshape(numpyArray1D, [2, 7, 2] as TensorShape)
+  expectEqual(TensorShape(reshaped3D.shape), [2, 7, 2])
+  let reshaped2D = np.reshape(reshaped3D, [14, 2] as TensorShape)
+  expectEqual(TensorShape(reshaped2D.shape), [14, 2])
+
   let numpyArrayStrided = np.array([[1, 2], [1, 2]], dtype: np.int32)[
       Python.slice(Python.None), 1]
   // Assert that the array has a stride, so that we're certainly testing a
@@ -95,6 +107,12 @@ NumpyConversionTests.test("tensor-conversion") {
                 tensor.array)
   }
 
+  let reshaped = np.reshape(numpyArrayInt32, [2, 3] as TensorShape)
+  if let tensor = expectNotNil(Tensor<Int32>(numpy: reshaped)) {
+    expectEqual(ShapedArray(shape: [2, 3], scalars: [1, 2, 3, 4, 5, 6]),
+                tensor.array)
+  }
+
   let numpyArrayStrided = np.array([[1, 2], [1, 2]], dtype: np.int32)[
       Python.slice(Python.None), 1]
   // Assert that the array has a stride, so that we're certainly testing a
@@ -118,6 +136,12 @@ NumpyConversionTests.test("tensor-round-trip") {
   let t3 = Tensor<Int32>(repeating: 30, shape: [8,5,4])
   expectEqual(t3, Tensor<Int32>(numpy: t3.makeNumpyArray())!)
 }
+
+NumpyConversionTests.test("tensor-shape") {
+  let pyArray = [2, 3].pythonObject
+  expectEqual(pyArray, TensorShape(2, 3).pythonObject)
+  expectEqual(TensorShape(2, 3), TensorShape(pyArray))
+}
 #endif
 
 runAllTests()