fix corner cases

app.py

@@ -16,6 +16,12 @@ t2 = torch.arange({n2}).view({dim2})
 
 """
 
+matrix_loop = """```python
+out = 0
+for i, j in zip(t1, t2):
+    out += i * j
+```
+"""
 
 def generate_example(dim1: list, dim2: list):
     n1 = 1
@@ -58,25 +64,35 @@ def sanitize_dimention(dim):
     return out
 
 
-def create_row(dim,is_dim=None,checks=None):
+def create_row(dim,is_dim=None,checks=None,version=1):
     out = "| "
     n_dim = len(dim)
     for i in range(n_dim):
-        # infered last dims
-        if (is_dim ==1 and i == n_dim-2) or (is_dim ==2 and i ==n_dim-1):
-            color = "green"
-            out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
-        # check every normal dimension
-        elif (is_dim ==1 and i != n_dim-1) or (is_dim ==2 and i ==n_dim-1):
-            color = "green" if checks[i] == "V" else "red"
-            out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
-        # checks last 2 dims
-        elif (is_dim ==1 and i == n_dim-1) or (is_dim ==2 and i ==n_dim-2):
-            color = "blue" if checks[i] == "V" else "yellow"
-            out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
-        # when using this function without checks
-        else : 
-            out+= f"{dim[i]} | "
+        if version == 1 :
+            # infered last dims
+            if (is_dim ==1 and i == n_dim-2) or (is_dim ==2 and i ==n_dim-1):
+                color = "green"
+                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
+            # check every normal dimension
+            elif (is_dim ==1 and i != n_dim-1) or (is_dim ==2 and i ==n_dim-1):
+                color = "green" if checks[i] == "V" else "red"
+                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
+            # checks last 2 dims
+            elif (is_dim ==1 and i == n_dim-1) or (is_dim ==2 and i ==n_dim-2):
+                color = "blue" if checks[i] == "V" else "yellow"
+                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
+            # when using this function without checks
+            else : 
+                out+= f"{dim[i]} | "
+        if version == 2 : 
+            if (is_dim == 1 and i != n_dim-1) : 
+                out += f"<strong style='color: green'> {dim[i]} </strong>| "
+            elif i == n_dim-1 :
+                color = "blue" if checks[i] == "V" else "yellow"
+                out += f"<strong style='color: {color}'> {dim[i]} </strong>| "
+            else : 
+                out += f"{dim[i]} | "
+            
     return out + "\n"
 
 
@@ -89,20 +105,20 @@ def create_header(n_dim, checks=None):
     return out
 
 
-def generate_table(dim1, dim2, checks=None):
+def generate_table(dim1, dim2, checks=None,version=1):
     n_dim = len(dim1)
     table = create_header(n_dim, checks)
     # tensor 1
     if not checks : 
         table += create_row(dim1)
     else : 
-        table += create_row(dim1,1,checks)
+        table += create_row(dim1,1,checks,version)
 
     # tensor 2
     if not checks : 
         table += create_row(dim2)
     else : 
-        table += create_row(dim2,2,checks)
+        table += create_row(dim2,2,checks,version)
     return table
 
 
@@ -122,8 +138,6 @@ def alignment_and_fill_with_ones(dim1, dim2):
     return dim1, dim2
 
 def check_validity(dim1,dim2):
-    if len(dim1) < 2:
-        return ["WIP"] * len(dim1)
     out = []
     for i in range(len(dim1)-2):
         if dim1[i] == dim2[i]:
@@ -138,8 +152,9 @@ def check_validity(dim1,dim2):
     return out
 
 
-def substitute_ones_with_concat(dim1,dim2):
-    for i in range(len(dim1)-1):
+def substitute_ones_with_concat(dim1,dim2,version = 1):
+    n = len(dim1)-2 if version ==1 else len(dim1)-1 
+    for i in range(n):
         dim1[i] = dim2[i] if dim1[i] == 1 else dim1[i]
         dim2[i] = dim1[i] if dim2[i] == 1 else dim2[i]
     return dim1, dim2
@@ -147,30 +162,63 @@ def substitute_ones_with_concat(dim1,dim2):
 def predict(dim1, dim2):
     dim1 = sanitize_dimention(dim1)
     dim2 = sanitize_dimention(dim2)
-    dim1, dim2, code = generate_example(dim1, dim2)
+    n1 , n2 = len(dim1) , len(dim2)
+    dim1, dim2, out = generate_example(dim1, dim2)
     # TODO 
-    # fix for dims if one or both have dimensions is 1
-    # Table 1 
-    dim1, dim2 = alignment_and_fill_with_ones(dim1, dim2)
-    table1 = generate_table(dim1, dim2)
-    # Table 2 
-    dim1, dim2 = substitute_ones_with_concat(dim1,dim2)
-    table2 = generate_table(dim1, dim2)
-    # Table 3 
-    checks = check_validity(dim1,dim2)
-    table3 = generate_table(dim1,dim2,checks)
-
-    out = code
-    out += "\n# Step1 (alignment and pre_append with ones)\n" + table1
-    out += "\n# Step2 (susbtitute columns that have 1 with concat)\nexcept for last 2 dimensions\n" + table2
-    out += "\n# Step3 (check if matrix multiplication is valid)\n" 
-    out += "* last dimension of dim1 should equal before last dimension of dim2 (blue or yellow colors)\n"
-    out += "* all the other dimensions should be equal to one another (green or red colors)\n\n" + table3
-    if "X" not in checks : 
-        dim1[-1] = dim2[-1]
-        out += "\n# Final dimension\n"
-        out+="as highlighted in <strong style='color:green'> green </strong> \n\n"
-        out+= f"`output.shape = {dim1}`"
+    if n1 >1 and n2 > 1 : 
+        # Table 1 
+        dim1, dim2 = alignment_and_fill_with_ones(dim1, dim2)
+        table1 = generate_table(dim1, dim2)
+        # Table 2 
+        dim1, dim2 = substitute_ones_with_concat(dim1,dim2)
+        table2 = generate_table(dim1, dim2)
+        # Table 3 
+        checks = check_validity(dim1,dim2)
+        table3 = generate_table(dim1,dim2,checks)
+
+        out += "\n# Step1 (alignment and pre_append with ones)\n" + table1
+        out += "\n# Step2 (susbtitute columns that have 1 with concat)\nexcept for last 2 dimensions\n" + table2
+        out += "\n# Step3 (check if matrix multiplication is valid)\n" 
+        out += "* last dimension of dim1 should equal before last dimension of dim2 (blue or yellow colors)\n"
+        out += "* all the other dimensions should be equal to one another (green or red colors)\n\n" + table3
+        if "X" not in checks : 
+            dim1[-1] = dim2[-1]
+            out += "\n# Final dimension\n"
+            out+="as highlighted in <strong style='color:green'> green </strong> \n\n"
+            out+= f"`output.shape = {dim1}`"
+    # case single dims
+    elif n1 == 1 and n2  == 1 :
+        out += "# Single Dimensional Cases\n"
+        out += "When both matricies have only single dims they should both have the same number of values in the first dimension\n"
+        out += "meaning that `t1.shape == t2.shape`\n"
+        out += "the output is a single value, think : \n"
+        out += matrix_loop
+    else : 
+        out += "# One of the dimensions is a single dimension\n"
+        out += "In this case we need to assert that the last dimension of `t1` "
+        out += "is equal to the last dimension of `t2`\n"
+        out += "Once the assertion is valid then we get rid of the last dimension and keep the rest\n"
+        out += "# Step 1 (alignment and fill with ones)\n"
+        dim1, dim2 = alignment_and_fill_with_ones(dim1, dim2)
+        table = generate_table(dim1, dim2)
+        out += table
+        out += "\n# Step2 (susbtitute columns that have 1 with concat)\n"
+        out += ""
+        dim1, dim2 = substitute_ones_with_concat(dim1,dim2,2)
+        checks = ["V"] * (len(dim1)-1)
+        if dim1[-1] == dim2[-1] :
+            checks.append("V")
+        else : 
+            checks.append("X")
+        table = generate_table(dim1, dim2,checks,2)
+        out+= table
+        if "X" not in checks : 
+            out += "\n#Final dimension"
+            out += "The final dimension is everything colored in <strong style='color:green'> green </strong> \n"
+            out += f"\nfinal dimension = `{dim1[:-1]}` "
+        
+        
+
     return out