From 6af94de00a2ffeaef8d09e4d46c2c50ef55b2e57 Mon Sep 17 00:00:00 2001
From: Michael Poutre <m1kep.my.mail@gmail.com>
Date: Wed, 6 Sep 2023 18:33:25 -0700
Subject: [PATCH] feat(func/max): Add function

---
 lib/__init__.py    |  2 ++
 lib/actions/max.py | 72 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 lib/actions/max.py

diff --git a/lib/__init__.py b/lib/__init__.py
index 5e463a7..6ff6fe8 100644
--- a/lib/__init__.py
+++ b/lib/__init__.py
@@ -1,5 +1,6 @@
 from custom_nodes.KepPromptLang.lib.actions.avg import AverageAction
 from custom_nodes.KepPromptLang.lib.actions.diff import DiffAction
+from custom_nodes.KepPromptLang.lib.actions.max import MaxAction
 from custom_nodes.KepPromptLang.lib.actions.mult import MultiplyAction
 from custom_nodes.KepPromptLang.lib.actions.neg import NegAction
 from custom_nodes.KepPromptLang.lib.actions.norm import NormAction
@@ -20,3 +21,4 @@ register_action(SlerpAction)
 register_action(AverageAction)
 register_action(ScaleDims)
 register_action(SetDims)
+register_action(MaxAction)
diff --git a/lib/actions/max.py b/lib/actions/max.py
new file mode 100644
index 0000000..04a177f
--- /dev/null
+++ b/lib/actions/max.py
@@ -0,0 +1,72 @@
+from typing import List, Union
+
+import torch
+from torch.nn import Embedding
+
+from custom_nodes.KepPromptLang.lib.action.base import Action, MultiArgAction
+from custom_nodes.KepPromptLang.lib.actions.action_utils import get_embedding
+from custom_nodes.KepPromptLang.lib.actions.utils import is_broadcastable
+from custom_nodes.KepPromptLang.lib.parser.prompt_segment import PromptSegment
+
+
+class MaxAction(MultiArgAction):
+    grammar = 'max(" arg ("|" arg)+ ")"'
+    name = "max"
+    chars = ["+", "+"]
+
+    def __init__(self, args: List[List[Union[PromptSegment, Action]]]) -> None:
+        super().__init__(args)
+
+        self.base_arg = args[0]
+        self.additional_args = args[1:]
+
+    def token_length(self) -> int:
+        # Max modifies the base segment, so the length is the length of the base segment
+        return sum(seg_or_action.token_length() for seg_or_action in self.base_arg)
+
+    def get_result(self, embedding_module: Embedding) -> torch.Tensor:
+        # Calculate the embeddings for the base segment
+        all_base_embeddings = [
+            get_embedding(seg_or_action, embedding_module)
+            for seg_or_action in self.base_arg
+        ]
+
+        result = torch.cat(all_base_embeddings, dim=1)
+
+        for arg in self.additional_args:
+            all_arg_embeddings = [
+                get_embedding(seg_or_action, embedding_module) for seg_or_action in arg
+            ]
+
+            arg_embedding = torch.cat(all_arg_embeddings, dim=1)
+
+            if is_broadcastable(result, arg_embedding):
+                result = torch.max(result, arg_embedding)
+            else:
+                print(
+                    "WARNING: shape mismatch when trying to apply max, arg will be averaged"
+                )
+                result = torch.max(result, torch.mean(arg_embedding, dim=1, keepdim=True))
+        return result
+
+    # def __repr__(self):
+    #     return f"sum(\n\tbase_segment={self.base_segment},\n\targs={self.args}\n)"
+    def __repr__(self) -> str:
+        return f"sum({', '.join(map(str, self.additional_args))})"
+
+    def depth_repr(self, depth=1):
+        out = "NudgeAction(\n"
+        if isinstance(self.base_arg, Action):
+            base_segment_repr = self.base_arg.depth_repr(depth + 1)
+            out += "\t" * depth + f"base_segment={base_segment_repr}\n"
+        else:
+            out += "\t" * depth + f"base_segment={self.base_arg.depth_repr()},\n"
+
+        if isinstance(self.additional_args, Action):
+            target_repr = self.additional_args.depth_repr(depth + 1)
+            out += "\t" * depth + f"target={target_repr},\n"
+        else:
+            out += "\t" * depth + f"target={self.additional_args.depth_repr()},\n"
+        out += "\t" * depth + f"weight={self.weight},\n"
+        out += "\t" * (depth - 1) + ")"
+        return out