Make 'to_standard' more lenient in terms of output shape

Author: asterycs

src/forward.jl

@@ -187,7 +187,15 @@ function evaluate(::Mult, arg1::BinaryOperation{Mult}, arg2::BinaryOperation{Mul
     new_args = []
 
     available1 = Any[arg1.arg1; arg1.arg2]
+    if is_elementwise_multiplication(arg1.arg1, arg1.arg2)
+        available1 = Any[arg1]
+    end
+
     available2 = Any[arg2.arg1; arg2.arg2]
+    if is_elementwise_multiplication(arg2.arg1, arg2.arg2)
+        available2 = Any[arg2]
+    end
+
 
     for i ∈ eachindex(available1)
         if isnothing(available1[i])
@@ -225,7 +233,7 @@ function evaluate(::Mult, arg1::BinaryOperation{Mult}, arg2::BinaryOperation{Mul
             if isnothing(new_arg)
                 return args[1]
             else
-                return evaluate(BinaryOperation{Mult}(new_arg, args[1]))
+                return BinaryOperation{Mult}(new_arg, args[1])
             end
         end
 

src/std.jl

@@ -629,6 +629,7 @@ end
 function to_standard(arg::BinaryOperation{Mult})
     target_indices = unique(get_free_indices(arg))
     target_len = length(target_indices)
+    is_scalar = isempty(target_indices)
 
     if length(target_indices) > 2
         throw_not_std(arg)
@@ -637,24 +638,18 @@ function to_standard(arg::BinaryOperation{Mult})
     l = to_standard(arg.arg1)
     r = to_standard(arg.arg2)
 
-    attempt = BinaryOperation{Mult}(l, r)
-    if get_free_indices(attempt) == target_indices
-        return attempt
-    end
-
-    attempt = BinaryOperation{Mult}(adjoint(l), r)
-    if get_free_indices(attempt) == target_indices
-        return attempt
-    end
-
-    attempt = BinaryOperation{Mult}(l, adjoint(r))
-    if get_free_indices(attempt) == target_indices
-        return attempt
-    end
-
-    attempt = BinaryOperation{Mult}(adjoint(l), adjoint(r))
-    if get_free_indices(attempt) == target_indices
-        return attempt
+    attempts = (
+        BinaryOperation{Mult}(l, r),
+        BinaryOperation{Mult}(adjoint(l), r),
+        BinaryOperation{Mult}(l, adjoint(r)),
+        BinaryOperation{Mult}(adjoint(l), adjoint(r)),
+    )
+
+    for attempt ∈ attempts
+        if length(get_free_indices(attempt)) == target_len &&
+           (is_scalar || last(get_free_indices(attempt)) == last(target_indices))
+            return attempt
+        end
     end
 
     throw_not_std(arg)