jaxable data

trunk/SUAVE/Core/Data.py

@@ -18,6 +18,8 @@
 
 from copy import copy
 
+from jax.tree_util import register_pytree_node_class
+
 # for enforcing attribute style access names
 import string
 from warnings import warn
@@ -33,6 +35,7 @@
 # ----------------------------------------------------------------------        
 
 ## @ingroup Core
+@register_pytree_node_class
 class Data(dict):
     """ An extension of the Python dict which allows for both tag and '.' usage.
         This is an unordered dictionary. So indexing it will not produce deterministic results.
@@ -45,6 +48,51 @@ class Data(dict):
         N/A
     """
 
+    def tree_flatten(self):
+        """ Returns a list of objects for tree operations.
+    
+            Assumptions:
+            N/A
+    
+            Source:
+            N/A
+    
+            Inputs:
+            method  - name of the method to access
+    
+            Outputs:
+            Result  - the results of the method
+    
+            Properties Used:
+            N/A    
+        """          
+        children = self.values() 
+        aux_data = None
+        return (children, aux_data)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        """ Rebuilds the Data class.
+    
+            Assumptions:
+            N/A
+    
+            Source:
+            N/A
+    
+            Inputs:
+            method  - name of the method to access
+    
+            Outputs:
+            Result  - the results of the method
+    
+            Properties Used:
+            N/A    
+        """          
+        return cls(*children)
+
+
+
     def __getattribute__(self, k):
         """ Retrieves an attribute set by a key k
     
@@ -851,4 +899,6 @@ def do_operation(A,B,C):
         # do the update!
         do_operation(self,other,result)    
 
-        return result
+        return result
+
+