Correctly pass the old config to the model sharding fiddler generation, which makes the output significantly less verbose. Also uses the new short namer in case intermediate variables are needed.

PiperOrigin-RevId: 549703159

Author: Fiddle-Config Team

fiddle/_src/codegen/codegen_diff.py

@@ -15,7 +15,9 @@
 
 """Library for converting generating fiddlers from diffs."""
 
+import abc
 import collections
+import dataclasses
 import functools
 import re
 import types
@@ -31,12 +33,63 @@
 import libcst as cst
 
 
+@dataclasses.dataclass(frozen=True)
+class ObjectToName:
+  prefix: str
+  path: daglish.Path
+
+  def __hash__(self):
+    return id(self)
+
+
+class VariableNamer(metaclass=abc.ABCMeta):
+  """Assigns names to intermediate variables."""
+
+  @abc.abstractmethod
+  def assign_names(self, all_to_name: List[ObjectToName]) -> List[str]:
+    """Assigns names, which will be further disambiguated by a Namespace."""
+
+
+class ExplicitNamer(VariableNamer):
+  """Assigns long and explicit names to all variables."""
+
+  def assign_names(self, all_to_name: List[ObjectToName]) -> List[str]:
+    return [
+        to_name.prefix + _path_to_name(to_name.path) for to_name in all_to_name
+    ]
+
+
+class ShorterNamer(VariableNamer):
+  """Assigns shorter, partly-disambiguated names to all variables."""
+
+  initial_suffix_elements: int = 1
+
+  def assign_names(self, all_to_name: List[ObjectToName]) -> List[str]:
+    name_to_paths = {}
+    for to_name in all_to_name:
+      sub_path = to_name.path[-self.initial_suffix_elements :]
+      name_to_paths.setdefault(
+          to_name.prefix + _path_to_name(sub_path), []
+      ).append(to_name)
+
+    result_as_dict = {}
+    for name, group in name_to_paths.items():
+      if len(group) == 1:
+        result_as_dict[group[0]] = name
+      else:
+        for to_name in group:
+          sub_path = to_name.path[-(self.initial_suffix_elements + 1) :]
+          result_as_dict[to_name] = to_name.prefix + _path_to_name(sub_path)
+    return [result_as_dict[to_name] for to_name in all_to_name]
+
+
 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,
+    variable_namer: Optional[VariableNamer] = None,
 ):
   """Returns the CST for a fiddler function that applies the changes in `diff`.
 
@@ -72,6 +125,8 @@ def fiddler_from_diff(
     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.
+    variable_namer: Object (basically just a function) to assign names to
+      intermediate variables.
 
   Returns:
     An `cst.Module` object.  You can convert this to a string using
@@ -90,25 +145,30 @@ def fiddler_from_diff(
     namespace.add(param_name)
     namespace.add(func_name)
 
+  variable_namer = ExplicitNamer() if variable_namer is None else variable_namer
+
   # Get a list of paths that are referenced by the diff.
   used_paths = _find_used_paths(diff)
 
   # Add variables for any used paths where the value (or any of the value's
   # ancestors) will be replaced by a change in the diff.  If we don't have an
   # `old` structure, then we pessimistically assume that we need to create
   # variables for all used paths.
-  moved_value_names = {}
+  moved_values_to_name = []
   if old is not None:
     modified_paths = set([change.target for change in diff.changes])
     _add_path_aliases(modified_paths, old)
     for path in sorted(used_paths, key=daglish.path_str):
       if any(path[:i] in modified_paths for i in range(len(path) + 1)):
-        moved_value_names[path] = namespace.get_new_name(
-            _path_to_name(path), f'moved_{param_name}_')
+        moved_values_to_name.append(ObjectToName(f'moved_{param_name}_', path))
   else:
     for path in sorted(used_paths, key=daglish.path_str):
-      moved_value_names[path] = namespace.get_new_name(
-          _path_to_name(path), f'original_{param_name}_')
+      moved_values_to_name.append(ObjectToName(f'original_{param_name}_', path))
+  initial_names = variable_namer.assign_names(moved_values_to_name)
+  moved_value_names = {
+      to_name.path: namespace.get_new_name(name, prefix='')
+      for to_name, name in zip(moved_values_to_name, initial_names)
+  }
 
   # Add variables for new shared values added by the diff.
   new_shared_value_names = [

fiddle/codegen/codegen_diff.py

@@ -16,4 +16,7 @@
 """Library for converting generating fiddlers from diffs."""
 
 # pylint: disable=unused-import
+from fiddle._src.codegen.codegen_diff import ExplicitNamer
 from fiddle._src.codegen.codegen_diff import fiddler_from_diff
+from fiddle._src.codegen.codegen_diff import ShorterNamer
+from fiddle._src.codegen.codegen_diff import VariableNamer