Turn 'Log' into a UnaryOperation

Author: asterycs

src/forward.jl

@@ -49,7 +49,7 @@ function diff(arg::UnaryOperation{Cos}, wrt::Variable)
     return BinaryOperation{Mult}(-UnaryOperation{Sin}(arg.arg), diff(arg.arg, wrt))
 end
 
-function diff(arg::Log, wrt::Variable)
+function diff(arg::UnaryOperation{Log}, wrt::Variable)
     return BinaryOperation{Mult}(
         BinaryOperation{Div}(Literal(1, get_free_indices(arg)...), arg.arg),
         diff(arg.arg, wrt),
@@ -413,20 +413,12 @@ function evaluate(::Mult, arg1::Zero, arg2::KrD)
     return Zero(free_indices...)
 end
 
-function evaluate(::Mult, arg1::UnaryOperation, arg2::KrD)
+function evaluate(::Mult, arg1::KrD, arg2::UnaryOperation)
     return evaluate(Mult(), arg2, arg1)
 end
 
-function evaluate(::Mult, arg1::KrD, arg2::UnaryOp) where {UnaryOp<:UnaryOperation}
-    if can_contract(evaluate(arg1), evaluate(arg2.arg))
-        return UnaryOp(evaluate(Mult(), evaluate(arg1), evaluate(arg2.arg)))
-    end
-
-    return BinaryOperation{Mult}(evaluate(arg2), evaluate(arg1))
-end
-
-function evaluate(::Mult, arg1::Log, arg2::KrD)
-    attempt = Log(evaluate(Mult(), evaluate(arg1.arg), evaluate(arg2)))
+function evaluate(::Mult, arg1::UnaryOperation{Op}, arg2::KrD) where {Op}
+    attempt = UnaryOperation{Op}(evaluate(Mult(), evaluate(arg1.arg), evaluate(arg2)))
 
     # This ensures that arg1.base and arg2 can contract and that the contraction is simple
     if length(get_free_indices(attempt)) == length(get_free_indices(arg1))
@@ -941,10 +933,6 @@ function _sub_from_product(arg1::BinaryOperation{Mult}, arg2::Value)
     return BinaryOperation{Sub}(evaluate(arg1), evaluate(arg2))
 end
 
-function evaluate(op::Log)
-    return Log(evaluate(op.arg))
-end
-
 function evaluate(op::Power)
     if op.exponent == 1
         return evaluate(op.base)

src/ir.jl

@@ -513,7 +513,7 @@ function to_ir(arg::BinaryOperation{Div})
     return ir.Quotient(to_ir(arg.arg1), to_ir(arg.arg2))
 end
 
-function to_ir(arg::Log)
+function to_ir(arg::UnaryOperation{Log})
     return ir.Log(to_ir(arg.arg))
 end
 

src/ricci.jl

@@ -109,12 +109,6 @@ end
 
 Base.hash(op::Power, h::UInt) = hash(op.exponent, hash(op.base, hash(Power, h)))
 
-struct Log <: Tensor
-    arg::Value
-end
-
-Base.hash(op::Log, h::UInt) = hash(op.arg, hash(Log, h))
-
 struct UnaryOperation{Op} <: Tensor where {Op}
     arg::Value
 end
@@ -125,6 +119,7 @@ struct Abs end
 struct Sgn end
 struct Sin end
 struct Cos end
+struct Log end
 
 function Base.sin(arg::Tensor)
     if !isempty(get_free_indices(arg))
@@ -243,10 +238,6 @@ function get_indices(arg::Power)
     return get_indices(arg.base)
 end
 
-function get_indices(arg::Log)
-    return get_indices(arg.arg)
-end
-
 function get_indices(arg::BinaryOperation{Op}) where {Op<:AdditiveOperation}
     arg1_free_ids, arg2_free_ids = get_free_indices.((arg.arg1, arg.arg2))
 
@@ -438,11 +429,11 @@ function Base.log(arg::Tensor)
         throw(DomainError(arg, "Argument is not a scalar, use log. for element-wise log."))
     end
 
-    return Log(arg)
+    return UnaryOperation{Log}(arg)
 end
 
 function Base.broadcasted(::typeof(log), arg::Tensor)
-    return Log(arg)
+    return UnaryOperation{Log}(arg)
 end
 
 function replace_letters(arg::BinaryOperation{Mult}, letter_map::Dict)
@@ -838,6 +829,10 @@ function to_string(arg::UnaryOperation{Cos})
     return "cos($(arg.arg))"
 end
 
+function to_string(arg::UnaryOperation{Log})
+    return "log($(arg.arg))"
+end
+
 function parenthesize(arg)
     return to_string(arg)
 end