Remove ir.PartialSum

Author: asterycs

src/ir.jl

@@ -95,11 +95,6 @@ struct Sum <: IR
     arg::IR
 end
 
-struct PartialSum <: IR
-    arg::IR
-    dim::Int
-end
-
 function _get_variables(arg::Mat)
     return _get_variables(arg.id)
 end
@@ -399,18 +394,6 @@ function to_ir(arg::BinaryOperation{Mult})
             else
                 return ir.Transpose(ir.Vec(ir.Literal(1)))
             end
-        elseif arg.arg1 isa Literal || arg.arg2 isa Literal
-            tensor = if arg.arg1 isa Literal
-                arg.arg2
-            else
-                arg.arg1
-            end
-
-            if typeof(target_indices[1]) == Upper
-                return ir.PartialSum(to_ir(tensor), 2)
-            else
-                return ir.PartialSum(to_ir(tensor), 1)
-            end
         end
     end
 

src/julia.jl

@@ -76,6 +76,12 @@ function to_julia(arg::ir.Quotient)
 end
 
 function to_julia(arg::ir.Product)
+    if arg.l isa ir.Transpose && arg.l.arg isa ir.Vec && arg.l.arg.id isa ir.Literal
+        return :($(to_julia(arg.l.arg.id)) .* (sum($(to_julia(arg.r)), dims = 1)))
+    elseif arg.r isa ir.Vec && arg.r.id isa ir.Literal
+        return :($(to_julia(arg.r.id)) .* (sum($(to_julia(arg.l)), dims = 2)))
+    end
+
     return :($(to_julia(arg.l)) * $(to_julia(arg.r)))
 end
 

src/stdstr.jl

@@ -200,15 +200,3 @@ end
 function to_std_str(arg::ir.Sum)
     return "sum(" * to_std_str(arg.arg) * ")"
 end
-
-function to_std_str(arg::ir.PartialSum)
-    if arg.dim == 1
-        product = ir.Product(ir.Transpose(ir.Vec(ir.Literal(1))), arg.arg) # Only used for dispatch
-        return "vec(1)ᵀ" * parenthesize(product, to_std_str, arg.arg)
-    elseif arg.dim == 2
-        product = ir.Product(arg.arg, ir.Vec(ir.Literal(1)))
-        return parenthesize(product, to_std_str, arg.arg) * "vec(1)"
-    end
-
-    throw(RuntimeError("Encountered a sum over an unsupported index"))
-end