fix: fix tensor grouping bug & optimize MMD calculation in causal_prediction/algorithms/regularization.py

dowhy/causal_prediction/algorithms/base_algorithm.py

@@ -26,8 +26,8 @@ def __init__(self, model, optimizer, lr, weight_decay, betas, momentum):
         self.momentum = momentum
 
         # Check if the optimizer is currently supported
-        if self.optimizer not in ["Adam", "SGD"]:
-            error_msg = self.optimizer + " is not implemented currently. Try Adam or SGD."
+        if self.optimizer not in ["Adam", "AdamW", "SGD"]:
+            error_msg = self.optimizer + " is not implemented currently. Try Adam, AdamW or SGD."
             raise Exception(error_msg)
 
     def training_step(self, train_batch, batch_idx):
@@ -89,6 +89,10 @@ def configure_optimizers(self):
             optimizer = torch.optim.Adam(
                 self.parameters(), lr=self.lr, weight_decay=self.weight_decay, betas=self.betas
             )
+        elif self.optimizer == "AdamW":
+            optimizer = torch.optim.AdamW(
+                self.parameters(), lr=self.lr, weight_decay=self.weight_decay, betas=self.betas
+            )
         elif self.optimizer == "SGD":
             optimizer = torch.optim.SGD(
                 self.parameters(), lr=self.lr, weight_decay=self.weight_decay, momentum=self.momentum

dowhy/causal_prediction/algorithms/cacm.py

@@ -16,9 +16,9 @@ def __init__(
         momentum=0.9,
         kernel_type="gaussian",
         ci_test="mmd",
-        attr_types=[],
+        attr_types=None,
         E_conditioned=True,
-        E_eq_A=[],
+        E_eq_A=None,
         gamma=1e-6,
         lambda_causal=1.0,
         lambda_conf=1.0,
@@ -97,13 +97,13 @@ def training_step(self, train_batch, batch_idx):
         objective /= nmb
         loss = objective
 
-        if self.attr_types != []:
+        if self.attr_types is not None:
             for attr_type_idx, attr_type in enumerate(self.attr_types):
                 attribute_labels = [
                     ai for _, _, ai in minibatches
                 ]  # [(batch_size, num_attrs)_1, batch_size, num_attrs)_2, ..., (batch_size, num_attrs)_(num_environments)]
 
-                E_eq_A_attr = attr_type_idx in self.E_eq_A
+                E_eq_A_attr = False if self.E_eq_A is None else attr_type_idx in self.E_eq_A
 
                 # Acause regularization
                 if attr_type == "causal":