@@ -90,14 +90,12 @@ class LazyTensorTraceBuilder {
90
90
}
91
91
92
92
/// Trace the given function and return the trace.
93
- static func trace< In: TensorGroup , Out: TensorGroup > (
94
- _ fn: ( In ) -> Out
95
- ) -> LazyTensorTrace {
96
- assert ( _RuntimeConfig. useLazyTensor, " Lazy tensor is not enabled for tracing. " )
93
+ static func trace< In: TensorGroup , Out: TensorGroup > ( _ fn: ( In ) -> Out ) -> LazyTensorTrace {
94
+ precondition ( _RuntimeConfig. useLazyTensor, " Lazy tensor is not enabled for tracing. " )
97
95
98
96
// Set up inputs for running `fn`
99
- let inputs = In . _typeList. map { Self . makePlaceholder ( with: $0) }
100
- let inputHandles = inputs . map { LazyTensorHandle ( _lazy: $0, index: 0 ) }
97
+ let inputOps = In . _typeList. map { Self . makePlaceholder ( with: $0) }
98
+ let inputHandles = inputOps . map { LazyTensorHandle ( _lazy: $0, index: 0 ) }
101
99
let input = In ( _handles: inputHandles)
102
100
103
101
// Run the function.
@@ -106,24 +104,21 @@ class LazyTensorTraceBuilder {
106
104
// Set up the closure that determines if a `LazyTensorOperation` should be an output.
107
105
let outputLazyOperations = output. _tensorHandles. map { ( handle: _AnyTensorHandle ) -> LazyTensorOperation in
108
106
let lazyOp = lazyTensorOperation ( handle)
109
- assert ( lazyOp != nil , " Found a non-lazy tensor in output when tracing. " )
107
+ precondition ( lazyOp != nil , " Found a non-lazy tensor in output when tracing. " )
110
108
return lazyOp!
111
109
}
112
- let outputIds = Set < ObjectIdentifier > ( outputLazyOperations. map {
113
- ObjectIdentifier ( $0)
114
- } )
110
+ let outputIds = Set < ObjectIdentifier > ( outputLazyOperations. map { ObjectIdentifier ( $0) } )
115
111
let isOutput : ( LazyTensorOperation ) -> Bool = { outputIds. contains ( ObjectIdentifier ( $0) ) }
116
112
117
113
// Create the builder and get the trace.
118
114
let builder = LazyTensorTraceBuilder ( )
119
115
builder. neverPromoteConstants = true
120
116
builder. isOutput = isOutput
121
- /// Set up the inputs for the builder as we need to have specific order.
122
- for inputOp in inputs {
123
- let id = ObjectIdentifier ( inputOp)
124
- builder. updateOperationAndCache ( id, inputOp)
117
+ /// Set up the inputs for the builder as we need to have them in a specific order.
118
+ for inputOp in inputOps {
119
+ builder. updateOperationAndCache ( ObjectIdentifier ( inputOp) , inputOp)
125
120
}
126
- builder. inputs = inputs
121
+ builder. inputs = inputOps
127
122
for lazyOp in outputLazyOperations { _ = builder. collectLazyOperation ( lazyOp) }
128
123
return LazyTensorTrace (
129
124
inputs: builder. inputs,
0 commit comments