Accessors for input and output eltypes

SciML · Nov 5, 2023 · 77a0a45 · 77a0a45
1 parent d31880e
commit 77a0a45
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/src/basic.jl b/src/basic.jl
@@ -633,11 +633,11 @@ function cache_self(L::ComposedOperator, u::AbstractVecOrMat)
         M = size(op, 1)
         sz = u isa AbstractMatrix ? (M, K) : (M,)
 
-        T = if op isa FunctionOperator # 
+        T = if op isa FunctionOperator #
             # FunctionOperator isn't guaranteed to play by the rules of
             # `promote_type`. For example, an irFFT is a complex operation
             # that accepts complex vector and returns  ones.
-            op.traits.eltypes[2]
+            output_eltype(op)
         else
             promote_type(eltype.((op, cache[1]))...)
         end

diff --git a/src/func.jl b/src/func.jl
@@ -100,6 +100,18 @@ function set_cache(f::FunctionOperator{iip, oop, mul5, T, F, Fa, Fi, Fai, Tr, P,
         f.p, f.t, cache)
 end
 
+function input_eltype(::FunctionOperator{iip, oop, mul5, T, F, Fa, Fi, Fai, Tr, P, Tt, C,
+        iType, oType,
+    }) where {iip, oop, mul5, T, F, Fa, Fi, Fai, Tr, P, Tt, C, iType, oType}
+    return iType
+end
+
+function output_eltype(::FunctionOperator{iip, oop, mul5, T, F, Fa, Fi, Fai, Tr, P, Tt, C,
+        iType, oType,
+    }) where {iip, oop, mul5, T, F, Fa, Fi, Fai, Tr, P, Tt, C, iType, oType}
+    return oType
+end
+
 """
 $(SIGNATURES)
 
@@ -301,9 +313,6 @@ function FunctionOperator(op,
         eltype(input), eltype(output)}(op,
         _op_adjoint, op_inverse, _op_adjoint_inverse, traits, p, _t, cache)
 
-    # L = FunctionOperator(op, _op_adjoint, op_inverse, _op_adjoint_inverse, traits, p, _t,
-    #     cache, eltype(input), eltype(output))
-
     # create cache
 
     if _unwrap_val(ifcache) & (L.cache === nothing)