diff --git a/stdlib/public/TensorFlow/TensorShape.swift b/stdlib/public/TensorFlow/TensorShape.swift index 7209e4d01a275..68a16136f30bd 100644 --- a/stdlib/public/TensorFlow/TensorShape.swift +++ b/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) + } + } +} diff --git a/test/Python/python_runtime.swift b/test/Python/python_runtime.swift index c44d2d9e60f14..5dbbc239aa4c5 100644 --- a/test/Python/python_runtime.swift +++ b/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(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") { diff --git a/test/TensorFlow/integration.swift b/test/TensorFlow/integration.swift index bb21daba09f79..ed1fd83063c1a 100644 --- a/test/TensorFlow/integration.swift +++ b/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") diff --git a/test/TensorFlowRuntime/numpy_conversion.swift b/test/TensorFlowRuntime/numpy_conversion.swift index 7788b08e85ca5..b27b7002d3009 100644 --- a/test/TensorFlowRuntime/numpy_conversion.swift +++ b/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(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(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(repeating: 30, shape: [8,5,4]) expectEqual(t3, Tensor(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()