Merge pull request #385 from MilesCranmer/reduce-allocs

refactor: reduce allocations from type assertion
MilesCranmer · Dec 11, 2024 · 5714142 · 5714142
2 parents ad13595 + 652ea0b
commit 5714142
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/InterfaceDynamicExpressions.jl b/src/InterfaceDynamicExpressions.jl
@@ -79,9 +79,8 @@ which speed up evaluation significantly.
 function expected_array_type(X::AbstractArray, ::Type)
     return typeof(similar(X, axes(X, 2)))
 end
-function expected_array_type(X::AbstractArray, ::Type, ::Val{:eval_grad_tree_array})
-    return typeof(X)
-end
+expected_array_type(X::AbstractArray, ::Type, ::Val{:eval_grad_tree_array}) = typeof(X)
+expected_array_type(::Matrix{T}, ::Type) where {T} = Vector{T}
 
 """
     eval_diff_tree_array(tree::Union{AbstractExpression,AbstractExpressionNode}, X::AbstractArray, options::AbstractOptions, direction::Int)

diff --git a/src/TemplateExpression.jl b/src/TemplateExpression.jl
@@ -372,7 +372,10 @@ end
         end
     end
 )
-@unstable IDE.expected_array_type(::AbstractMatrix, ::Type{<:TemplateExpression}) = Any
+@unstable begin
+    IDE.expected_array_type(::AbstractArray, ::Type{<:TemplateExpression}) = Any
+    IDE.expected_array_type(::Matrix{T}, ::Type{<:TemplateExpression}) where {T} = Any
+end
 
 function DA.violates_dimensional_constraints(
     @nospecialize(tree::TemplateExpression),