| 14 | from functools import partial |
| 15 | |
| 16 | class Variable(BaseModel): |
| 17 | model_config = ConfigDict(arbitrary_types_allowed=True) |
| 18 | |
| 19 | name: str = Field(description="The name of the variable.") |
| 20 | type: str = Field(description="The type of the variable.") |
| 21 | description: str = Field(description="The description of the variable.") |
| 22 | require_grad: bool = Field(default=False, description="Whether the variable requires gradient.") |
| 23 | template: Optional[str] = Field(default=None, description="The template of the variable.") |
| 24 | variables: Optional[Union[Dict[str, 'Variable'], 'Variable', Any]] = Field(default=None, description="The elements of the variable. Can be a dict (keyed by name), single Variable, or direct value.") |
| 25 | |
| 26 | # gradient related attributes |
| 27 | gradients: Set['Variable'] = Field(default_factory=set, description="Text gradients for this variable.") |
| 28 | gradients_context: Dict['Variable', str] = Field(default_factory=lambda: defaultdict(lambda: None), description="Context for gradients.") |
| 29 | grad_fn: Optional[Any] = Field(default=None, description="Gradient function for backward pass.") |
| 30 | predecessors: Set['Variable'] = Field(default_factory=set, description="Predecessor variables in computation graph.") |
| 31 | reduce_meta: List[Dict] = Field(default_factory=list, description="Metadata for gradient reduction.") |
| 32 | |
| 33 | def __hash__(self): |
| 34 | return id(self) |
| 35 | |
| 36 | def __eq__(self, other): |
| 37 | return id(self) == id(other) |
| 38 | |
| 39 | @classmethod |
| 40 | def from_dict(cls, data: Dict[str, Any]) -> 'Variable': |
| 41 | """Recursively construct Variable tree from nested dict. |
| 42 | |
| 43 | Supports dict format for variables: |
| 44 | - Dict format: {"var1": {"name": "var1", ...}, "var2": {"name": "var2", ...}} |
| 45 | """ |
| 46 | subvars = data.get("variables") |
| 47 | if isinstance(subvars, dict): |
| 48 | # Dict format: recursively process each variable |
| 49 | subvars = {k: cls.from_dict(v) if isinstance(v, dict) and "name" in v else v |
| 50 | for k, v in subvars.items()} |
| 51 | elif subvars is not None and not isinstance(subvars, dict): |
| 52 | # Direct value (string, etc.) - keep as is |
| 53 | pass |
| 54 | return cls( |
| 55 | name=data["name"], |
| 56 | type=data.get("type", ""), |
| 57 | description=data.get("description", ""), |
| 58 | require_grad=data.get("require_grad", False), |
| 59 | template=data.get("template"), |
| 60 | variables=subvars, |
| 61 | ) |
| 62 | |
| 63 | def render(self, modules: Dict[str, Any]) -> str: |
| 64 | """Render the template with the given modules.""" |
| 65 | if self.template is None: |
| 66 | return "" |
| 67 | |
| 68 | env = Environment() |
| 69 | ast = env.parse(self.template) |
| 70 | vars_used = meta.find_undeclared_variables(ast) |
| 71 | ctx = dict(modules) |
| 72 | |
| 73 | for var in vars_used: |
no outgoing calls