@@ -9,6 +9,7 @@ use std::str::FromStr;
99use crate :: expand:: { Decodable , Encodable , HashStable_Generic } ;
1010use crate :: ptr:: P ;
1111use crate :: { Ty , TyKind } ;
12+ use crate :: expand:: typetree:: TypeTree ;
1213
1314/// Forward and Reverse Mode are well known names for automatic differentiation implementations.
1415/// Enzyme does support both, but with different semantics, see DiffActivity. The First variants
@@ -85,6 +86,9 @@ pub struct AutoDiffItem {
8586 /// The name of the function being generated
8687 pub target : String ,
8788 pub attrs : AutoDiffAttrs ,
89+ // --- TypeTree support ---
90+ pub inputs : Vec < TypeTree > ,
91+ pub output : TypeTree ,
8892}
8993
9094#[ derive( Clone , Eq , PartialEq , Encodable , Decodable , Debug , HashStable_Generic ) ]
@@ -112,6 +116,10 @@ impl AutoDiffAttrs {
112116 pub fn has_primal_ret ( & self ) -> bool {
113117 matches ! ( self . ret_activity, DiffActivity :: Active | DiffActivity :: Dual )
114118 }
119+ /// New constructor for type tree support
120+ pub fn into_item ( self , source : String , target : String , inputs : Vec < TypeTree > , output : TypeTree ) -> AutoDiffItem {
121+ AutoDiffItem { source, target, attrs : self , inputs, output }
122+ }
115123}
116124
117125impl DiffMode {
@@ -284,6 +292,8 @@ impl AutoDiffAttrs {
284292impl fmt:: Display for AutoDiffItem {
285293 fn fmt ( & self , f : & mut fmt:: Formatter < ' _ > ) -> fmt:: Result {
286294 write ! ( f, "Differentiating {} -> {}" , self . source, self . target) ?;
287- write ! ( f, " with attributes: {:?}" , self . attrs)
295+ write ! ( f, " with attributes: {:?}" , self . attrs) ?;
296+ write ! ( f, " with inputs: {:?}" , self . inputs) ?;
297+ write ! ( f, " with output: {:?}" , self . output)
288298 }
289299}
0 commit comments