Robustify the Branch to allow branches to be None (using the logical and)

Author: YoanSallami

coverage-badge.svg

@@ -148,38 +148,18 @@ async def call(self, inputs, training=False):
         )
         choice = decision.get("choice")
         outputs = []
+        if self.inject_decision:
+            inputs = await ops.concat(
+                inputs,
+                decision,
+                name=self.name + "_inputs_with_decision",
+            )
         for label, module in self.branches.items():
             if label == choice:
                 if module:
-                    if self.inject_decision and self.return_decision:
+                    if self.return_decision:
                         outputs.append(
-                            await ops.concat(
-                                decision,
-                                await module(
-                                    await ops.concat(
-                                        inputs,
-                                        decision,
-                                        name=self.name + "_inputs_with_decision",
-                                    ),
-                                    training=training,
-                                ),
-                                name=self.name + "_with_decision",
-                            )
-                        )
-                    elif self.inject_decision and not self.return_decision:
-                        outputs.append(
-                            await module(
-                                await ops.concat(
-                                    inputs,
-                                    decision,
-                                    name=self.name + "_inputs_with_decision",
-                                ),
-                                training=training,
-                            )
-                        )
-                    elif not self.inject_decision and self.return_decision:
-                        outputs.append(
-                            await ops.concat(
+                            await ops.logical_and(
                                 decision,
                                 await module(
                                     inputs,
@@ -207,36 +187,16 @@ async def compute_output_spec(self, inputs, training=False):
             inputs,
             training=training,
         )
+        if self.inject_decision:
+            inputs = await ops.concat(
+                inputs,
+                decision,
+                name=self.name + "_inputs_with_decision",
+            )
         for module in self.branches.values():
-            if self.inject_decision and self.return_decision:
-                outputs.append(
-                    await ops.concat(
-                        decision,
-                        await module(
-                            await ops.concat(
-                                inputs,
-                                decision,
-                                name=self.name + "_inputs_with_decision",
-                            ),
-                            training=training,
-                        ),
-                        name=self.name + "_with_decision",
-                    )
-                )
-            elif self.inject_decision and not self.return_decision:
-                outputs.append(
-                    await module(
-                        await ops.concat(
-                            inputs,
-                            decision,
-                            name=self.name + "_inputs_with_decision",
-                        ),
-                        training=training,
-                    )
-                )
-            elif not self.inject_decision and self.return_decision:
+            if self.return_decision:
                 outputs.append(
-                    await ops.concat(
+                    await ops.logical_and(
                         decision,
                         await module(
                             inputs,

synalinks/src/modules/core/branch.py

@@ -3,7 +3,7 @@
 from synalinks.src.api_export import synalinks_export
 
 # Unique source of truth for the version number.
-__version__ = "0.2.023"
+__version__ = "0.2.024"
 
 
 @synalinks_export("synalinks.version")