Skip to content
This repository was archived by the owner on Jul 1, 2023. It is now read-only.

Commit baa8cd4

Browse files
committed
assert->precondition, misc cleanup.
1 parent f4d32f7 commit baa8cd4

File tree

1 file changed

+10
-15
lines changed

1 file changed

+10
-15
lines changed

Sources/TensorFlow/Core/LazyTensorTrace.swift

Lines changed: 10 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -90,14 +90,12 @@ class LazyTensorTraceBuilder {
9090
}
9191

9292
/// 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.")
9795

9896
// 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) }
10199
let input = In(_handles: inputHandles)
102100

103101
// Run the function.
@@ -106,24 +104,21 @@ class LazyTensorTraceBuilder {
106104
// Set up the closure that determines if a `LazyTensorOperation` should be an output.
107105
let outputLazyOperations = output._tensorHandles.map { (handle: _AnyTensorHandle) -> LazyTensorOperation in
108106
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.")
110108
return lazyOp!
111109
}
112-
let outputIds = Set<ObjectIdentifier>(outputLazyOperations.map {
113-
ObjectIdentifier($0)
114-
})
110+
let outputIds = Set<ObjectIdentifier>(outputLazyOperations.map { ObjectIdentifier($0) })
115111
let isOutput: (LazyTensorOperation) -> Bool = { outputIds.contains(ObjectIdentifier($0)) }
116112

117113
// Create the builder and get the trace.
118114
let builder = LazyTensorTraceBuilder()
119115
builder.neverPromoteConstants = true
120116
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)
125120
}
126-
builder.inputs = inputs
121+
builder.inputs = inputOps
127122
for lazyOp in outputLazyOperations { _ = builder.collectLazyOperation(lazyOp) }
128123
return LazyTensorTrace(
129124
inputs: builder.inputs,

0 commit comments

Comments
 (0)