Pax-specific flattened config codegen (MVP, will be building on it further soon).

This will help flatten nested experiment hierarchies, making configs for baseline models easier to read.

PiperOrigin-RevId: 543513463

Author: Fiddle-Config Team

fiddle/_src/codegen/auto_config/code_ir.py

@@ -121,6 +121,9 @@ class AttributeExpression(CodegenNode):
   base: Any  # Wrapped expression, can involve VariableReference's
   attribute: str
 
+  def __hash__(self):
+    return id(self)
+
 
 @dataclasses.dataclass
 class ArgFactoryExpr(CodegenNode):

fiddle/_src/codegen/auto_config/ir_printer.py

@@ -90,6 +90,24 @@ def traverse(value, state: daglish.State) -> str:
     elif isinstance(value, code_ir.WithTagsCall):
       sub_value = state.map_children(value).expression
       return f"WithTagsCall[{sub_value}]"
+    elif isinstance(value, code_ir.SymbolOrFixtureCall):
+      symbol_expression = state.call(
+          value.symbol_expression, daglish.Attr("symbol_expression")
+      )
+      positional_arg_expressions = state.call(
+          value.positional_arg_expressions,
+          daglish.Attr("positional_arg_expressions"),
+      )
+      arg_expressions = state.call(
+          value.arg_expressions, daglish.Attr("arg_expressions")
+      )
+      return (
+          f"call:<{symbol_expression}"
+          f"(*[{positional_arg_expressions}],"
+          f" **{arg_expressions})>"
+      )
+    elif isinstance(value, code_ir.ModuleReference):
+      return value.name.value
     elif isinstance(value, code_ir.Name):
       return value.value
     elif isinstance(value, type):

fiddle/_src/codegen/auto_config/ir_to_cst.py

@@ -154,6 +154,7 @@ def _prepare_args_helper(
       except:
         print(f"\n\nERROR CONVERTING: {value!r}")
         print(f"\n\nTYPE: {type(value)}")
+        print("Path", daglish.path_str(state.current_path))
         raise
 
   return daglish.MemoizedTraversal.run(traverse, expr)
@@ -198,7 +199,7 @@ def code_for_fn(
           ),
       ]
   )
-  if fn.parameters:
+  if fn.parameters and len(fn.parameters) > 1:
     whitespace_before_params = cst.ParenthesizedWhitespace(
         cst.TrailingWhitespace(),
         indent=True,

fiddle/_src/codegen/codegen_diff.py

@@ -19,7 +19,7 @@
 import functools
 import re
 import types
-from typing import Any, Callable, Dict, List, Set, Tuple
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple
 
 from fiddle import daglish
 from fiddle import diffing
@@ -31,10 +31,13 @@
 import libcst as cst
 
 
-def fiddler_from_diff(diff: diffing.Diff,
-                      old: Any = None,
-                      func_name: str = 'fiddler',
-                      param_name: str = 'cfg'):
+def fiddler_from_diff(
+    diff: diffing.Diff,
+    old: Any = None,
+    func_name: str = 'fiddler',
+    param_name: str = 'cfg',
+    import_manager: Optional[import_manager_lib.ImportManager] = None,
+):
   """Returns the CST for a fiddler function that applies the changes in `diff`.
 
   The returned `cst.Module` consists of a set of `import` statements for any
@@ -66,18 +69,26 @@ def fiddler_from_diff(diff: diffing.Diff,
       all referenced paths.
     func_name: The name for the fiddler function.
     param_name: The name for the parameter to the fiddler function.
+    import_manager: Existing import manager. Usually set to None, but if you are
+      integrating this with other code generation tasks, it can be nice to
+      share.
 
   Returns:
     An `cst.Module` object.  You can convert this to a string using
     `result.code`.
   """
-  # Create a namespace to keep track of variables that we add.  Reserve the
-  # names of the param & func.
-  namespace = namespace_lib.Namespace()
-  namespace.add(param_name)
-  namespace.add(func_name)
-
-  import_manager = import_manager_lib.ImportManager(namespace)
+  if import_manager is None:
+    # Create a namespace to keep track of variables that we add.  Reserve the
+    # names of the param & func.
+    namespace = namespace_lib.Namespace()
+    namespace.add(param_name)
+    namespace.add(func_name)
+
+    import_manager = import_manager_lib.ImportManager(namespace)
+  else:
+    namespace = import_manager.namespace
+    namespace.add(param_name)
+    namespace.add(func_name)
 
   # Get a list of paths that are referenced by the diff.
   used_paths = _find_used_paths(diff)