@@ -563,7 +563,7 @@ private void _add_object_graph_edges(SavedObject proto, int node_id)
563
563
return proto . KindCase switch
564
564
{
565
565
SavedObject . KindOneofCase . UserObject => _recreate_user_object ( proto . UserObject , node_id ) ,
566
- SavedObject . KindOneofCase . Function => _recreate_function ( proto . Function , null ) ,
566
+ SavedObject . KindOneofCase . Function => _recreate_function ( proto . Function , dependencies ) ,
567
567
SavedObject . KindOneofCase . BareConcreteFunction => _recreate_bare_concrete_function ( proto . BareConcreteFunction , dependencies ) ,
568
568
SavedObject . KindOneofCase . Variable => _recreate_variable ( proto . Variable ) ,
569
569
SavedObject . KindOneofCase . CapturedTensor => throw new NotImplementedException ( ) ,
@@ -626,7 +626,7 @@ private void _add_object_graph_edges(SavedObject proto, int node_id)
626
626
}
627
627
628
628
private ( Function , Action < object , object , object > ) _recreate_function ( SavedFunction proto ,
629
- Dictionary < OneOf < string , int > , Trackable > dependencies )
629
+ IDictionary < OneOf < string , int > , Trackable > dependencies )
630
630
{
631
631
var fn = function_deserialization . recreate_function ( proto , _concrete_functions ) ;
632
632
foreach ( var name in proto . ConcreteFunctions )
@@ -644,6 +644,13 @@ private void _add_object_graph_edges(SavedObject proto, int node_id)
644
644
return ( fn , setattr ) ;
645
645
}
646
646
647
+ private ( Tensor , Action < object , object , object > ) _get_tensor_from_fn ( CapturedTensor proto )
648
+ {
649
+ var outer_graph = _concrete_functions [ proto . ConcreteFunction ] . func_graph ;
650
+ var captured_tensor = outer_graph . get_tensor_by_name ( proto . Name ) ;
651
+ return ( captured_tensor , setattr ) ;
652
+ }
653
+
647
654
// TODO: remove this to a common class.
648
655
public static Action < object , object , object > setattr = ( x , y , z ) =>
649
656
{
0 commit comments