@@ -737,7 +737,7 @@ extension LazyTensorOperation {
737
737
// Return materialized outputs if any.
738
738
if let outputs = outputs { return outputs }
739
739
740
- LazyTensorOperation . materializeLiveTensors ( self )
740
+ materializeLiveTensors ( )
741
741
742
742
// Our outputs should have been updated by now. Otherwise,
743
743
// something terrible happened!
@@ -748,29 +748,35 @@ extension LazyTensorOperation {
748
748
/// Converts symbolic tensor inputs to concrete inputs if the
749
749
/// associated `LazyTensorOperation` has been materialized.
750
750
private func maybeMaterializeInputs( ) {
751
+ /// If `lazyTensor` is symbolic and the associated `LazyTensorOperation`
752
+ /// has been materialized, return the corresponding concrete `LazyTensor`.
753
+ /// Otherwise, return `lazyTensor` untouched.
751
754
func maybeMaterialized( lazyTensor: LazyTensor ) -> LazyTensor {
752
755
let handle = lazyTensor. handle
753
- if case let LazyTensor . Handle . symbolic( lazyOp, index, _) = handle {
754
- if let outputs = lazyOp. outputs {
755
- return LazyTensor ( _materialized: outputs [ index] )
756
- }
756
+ if case let LazyTensor . Handle . symbolic( lazyOp, index, _) = handle,
757
+ let outputs = lazyOp. outputs {
758
+ return LazyTensor ( _materialized: outputs [ index] )
757
759
}
758
760
return lazyTensor
759
761
}
760
762
763
+ /// Rewrite the input such that all symbolic values that have been
764
+ /// materialized have been replaced by the corresponding concerete
765
+ /// inputs. If no symbolic values have been materialized or if there are
766
+ /// no symbolic values, return the `input` untouched.
761
767
func maybeMaterialized( input: Input ) -> Input {
762
768
switch input {
763
769
case . single( let h) :
764
- return Input . single ( maybeMaterialized ( lazyTensor: h) )
770
+ return . single( maybeMaterialized ( lazyTensor: h) )
765
771
case . list( let elements) :
766
- return Input . list ( elements. map { maybeMaterialized ( lazyTensor: $0) } )
772
+ return . list( elements. map { maybeMaterialized ( lazyTensor: $0) } )
767
773
}
768
774
}
769
775
inputs = inputs. map { maybeMaterialized ( input: $0) }
770
776
}
771
777
772
- private static func materializeLiveTensors( _ lazyOp : LazyTensorOperation ) {
773
- let lazyTrace = LazyTensorTrace ( lazyOp )
778
+ private func materializeLiveTensors( ) {
779
+ let lazyTrace = LazyTensorTrace ( self )
774
780
debugLog ( " Extracted trace: \n \( lazyTrace) " )
775
781
776
782
let function = TFFunction ( trace: lazyTrace)
0 commit comments