public SubGraph getSubGraph ( SameDiff sd , DifferentialFunction rootFn ) { Preconditions . checkState ( matches ( sd , rootFn ) , "Root function does not match predicate" ) ; List < DifferentialFunction > childNodes = new ArrayList <> ( ) ; //Need to work out child nodes if ( ! opInputSubgraphPredicates . isEmpty ( ) ) { for ( Map . Entry < Integer , OpPredicate > entry : opInputSubgraphPredicates . entrySet ( ) ) { OpPredicate p2 = entry . getValue ( ) ; SDVariable arg = rootFn . arg ( entry . getKey ( ) ) ; DifferentialFunction df = sd . getVariableOutputFunction ( arg . getVarName ( ) ) ; if ( df != null ) { childNodes . add ( df ) ; if ( p2 instanceof SubGraphPredicate ) { SubGraph sg = ( ( SubGraphPredicate ) p2 ) . getSubGraph ( sd , df ) ; childNodes . addAll ( sg . childNodes ) ; } } } } SubGraph sg = SubGraph . builder ( ) . sameDiff ( sd ) . rootNode ( rootFn ) . childNodes ( childNodes ) . build ( ) ; return sg ; }
Get the SubGraph that matches the predicate