High-quality tensor operation preconditions in Swift

[ematejska]

Imported from JIRA, filed by Dan Zheng (https://bugs.swift.org/browse/TF-655)

Add high-quality tensor operation preconditions in Swift instead of relying on TensorFlow.

Currently, TensorFlow runtime errors are triggered, with an ugly stack trace:

import TensorFlow
let x = Tensor<Float>(ones: [3])
print(matmul(x, x))
$ swift matmul.swift
Fatal error: In[0] is not a matrix. Instead it has shape [3]: file /Users/danielzheng/swift-tf/tensorflow-swift-apis/Sources/TensorFlow/Bindings/EagerExecution.swift, line 299
Stack dump:
0.	Program arguments: /Library/Developer/Toolchains/swift-tensorflow-RELEASE-0.4.xctoolchain/usr/bin/swift -frontend -interpret matmul.swift -enable-objc-interop -sdk /Applications/Xcode-beta.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX10.15.sdk -color-diagnostics -module-name matmul
1.	Swift version 5.1-dev (LLVM af1f73e9e9, Swift 2761ac9f87)
0  swift                    0x0000000111c226d5 llvm::sys::PrintStackTrace(llvm::raw_ostream&) + 37
1  swift                    0x0000000111c21718 llvm::sys::RunSignalHandlers() + 248
2  swift                    0x0000000111c22cc8 SignalHandler(int) + 264
3  libsystem_platform.dylib 0x00007fff5bec9b5d _sigtramp + 29
4  libsystem_platform.dylib 0x00007ffee1c550f0 _sigtramp + 2245572016
5  libswiftCore.dylib       0x000000011960d8d9 $ss17_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_SSAHSus6UInt32VtF + 25
6  libswiftTensorFlow.dylib 0x0000000119e2c01c $s10TensorFlow24internalConsistencyCheck__4file4lineySb_SSs12StaticStringVSutF + 316
7  libswiftTensorFlow.dylib 0x0000000119c927ae $s10TensorFlow7checkOk_4file4lineys13OpaquePointerVSg_s12StaticStringVSutF + 462
8  libswiftTensorFlow.dylib 0x0000000119c9992a $s10TensorFlow6TFE_OpV14evaluateUnsafeSpys13OpaquePointerVGyF + 506
9  libswiftTensorFlow.dylib 0x0000000119c9a024 $s10TensorFlow6TFE_OpV7executeyxSiAA0A13ArrayProtocolRzlF + 132
10 libswiftTensorFlow.dylib 0x0000000119ca2c34 $s10TensorFlow6TFE_OpVAA17TFTensorOperationA2aDP7executeyqd__SiAA0A13ArrayProtocolRd__lFTW + 52
11 libswiftTensorFlow.dylib 0x0000000119d36e65 $s10TensorFlow3RawO6matMul__10transposeA0F1BAA0A0VyxGAI_AIS2btSjRzAA0aB6ScalarRzlFZ + 1221
12 libswiftTensorFlow.dylib 0x0000000119e9f263 $s10TensorFlow6matmul_10transposed_AcA0A0VyxGAF_SbAFSbtSjRzAA0aB6ScalarRzlF + 1427
13 libswiftTensorFlow.dylib 0x000000011a3b7109 $s10TensorFlow6matmul_10transposed_AcA0A0VyxGAF_SbAFSbtSjRzAA0aB6ScalarRzlF + 5342265
14 swift                    0x000000010e01ed58 llvm::MCJIT::runFunction(llvm::Function*, llvm::ArrayRef<llvm::GenericValue>) + 456
15 swift                    0x000000010e021ec1 llvm::ExecutionEngine::runFunctionAsMain(llvm::Function*, std::__1::vector<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> >, std::__1::allocator<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > > const&, char const* const*) + 1313
16 swift                    0x000000010e0144cf swift::RunImmediately(swift::CompilerInstance&, std::__1::vector<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> >, std::__1::allocator<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > > const&, swift::IRGenOptions&, swift::SILOptions const&) + 3455
17 swift                    0x000000010e004ce9 performCompile(swift::CompilerInstance&, swift::CompilerInvocation&, llvm::ArrayRef<char const*>, int&, swift::FrontendObserver*, swift::UnifiedStatsReporter*) + 12985
18 swift                    0x000000010e000a51 swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 3025
19 swift                    0x000000010dfa8e59 main + 729
20 libdyld.dylib            0x00007fff5bcde3d5 start + 1
[1]    4069 illegal hardware instruction  swift matmul.swift

PyTorch is a good source of inspiration:

import torch

x = torch.ones([3, 4, 5, 6])
torch.matmul(x, x)
Traceback (most recent call last):
  File "matmul.py", line 4, in <module>
    torch.matmul(x, x)
RuntimeError: Expected tensor to have size 6 at dimension 1, but got size 5 for argument #2 'batch2' (while checking arguments for bmm)

• Comment by Anthony Platanios [ 17 Jul 2019 ]

This would be great, but using the current implementation of Tensor.shape and Tensor.rank it would cause lazy tensor materialization. We would need to first add support for shape inference without requiring materialization of lazy tensors.Also, is there no way to provide a more information stack trace whenever TensorFlow runtime errors are triggered?

• Comment by Anthony Platanios [ 17 Jul 2019 ]
  Currently S4TF only shows the last step of the stack trace which is always in EagerExecution.swift. Showing a more complete stack trace would go a long way helping with this.
• Comment by Dan Zheng [ 17 Jul 2019 ]

I believe Gogul Balakrishnan is working on this (shape propagation for LazyTensor).
Though tensor operation preconditions (when addeded) will catch errors early so that ugly TensorFlow runtime stack traces don't occur.Here's a related recent PR improving Swift runtime failure messages (involves SIL and DebugInfo support): swiftlang/swift#25978

[dan-zheng]

Thanks for importing!

I believe this is a duplicate of #517, which already has some discussion. Let's use that issue instead!