account for both cases when chemsys: str | list[str]

Author: andrew31416

mattergen/common/data/collate.py

@@ -237,20 +237,8 @@ def _merge(xs: list[PyTree[T]], structure: PyTree[int]) -> PyTree[T]:
                     )
                     del x[attr]  # type: ignore
 
-        # Batch.from_data_list will concat attr: list[list] to list[list[list]], we need to handle separately
-        attr_is_twod_list = []
-
-        for attr in attrs:
-            if all(isinstance(x[attr], list) for x in xs) and all(isinstance(_x,list) for x in xs for _x in x[attr]):
-                attr_is_twod_list.append(attr)
-
         try:
             batch = Batch.from_data_list(xs)
-
-            # handle attr: list[list] as a special case
-            for attr in attr_is_twod_list:
-                # convert batch.attr: list[list[list]] to list[list]
-                batch[attr] = list(chain(*[x[attr] for x in xs]))
         except Exception as e:
             # Check if dtypes do not match:
             for attr in attrs:

mattergen/common/data/types.py

@@ -11,7 +11,7 @@
 from mattergen.common.utils.globals import PROPERTY_SOURCE_IDS
 
 PropertySourceId = str
-TargetProperty = dict[PropertySourceId, int | float | Sequence[str]]
+TargetProperty = dict[PropertySourceId, int | float | str | Sequence[str]]
 
 
 @dataclass(frozen=True)

mattergen/diffusion/sampling/classifier_free_guidance.py

@@ -75,10 +75,18 @@ def get_conditional_score():
             batch_no_condition = self._remove_conditioning_fn(x)
             batch_with_condition = self._keep_conditioning_fn(x)
             joint_batch = collate([batch_no_condition, batch_with_condition])
+
+            for attr,value in batch_no_condition.items():
+                if isinstance(value, list):
+                    joint_batch[attr] = batch_no_condition[attr]+batch_with_condition[attr]
+
+
             combined_score = super(GuidedPredictorCorrector, self)._score_fn(
                 x=joint_batch, t=torch.cat([t, t], dim=0),
             )
             # Split the combined score back into unconditional and conditional parts.
+            # Any batch.attr: list fields will be wrong here because of the manual concatenation above
+            # this should be ok as self._multi_corruption.corrupted_fields are always torch.Tensor
             unconditional_score = combined_score[0]
             conditional_score = combined_score[1]