Merge pull request #89 from SciML/type

allow for dispatching get_tmp on the type
SciML · Dec 31, 2023 · e379918 · e379918
2 parents 6636083 + c5a7702
commit e379918
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "PreallocationTools"
 uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "0.4.13"
+version = "0.4.14"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/PreallocationTools.jl b/src/PreallocationTools.jl
@@ -46,6 +46,15 @@ function get_tmp(dc::FixedSizeDiffCache, u::T) where {T <: ForwardDiff.Dual}
     end
 end
 
+function get_tmp(dc::FixedSizeDiffCache, u::Type{T}) where {T <: ForwardDiff.Dual}
+    x = reinterpret(T, dc.dual_du)
+    if chunksize(T) === chunksize(eltype(dc.dual_du))
+        x
+    else
+        @view x[axes(dc.du)...]
+    end
+end
+
 function get_tmp(dc::FixedSizeDiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual}
     x = reinterpret(T, dc.dual_du)
     if chunksize(T) === chunksize(eltype(dc.dual_du))
@@ -116,6 +125,14 @@ function get_tmp(dc::DiffCache, u::T) where {T <: ForwardDiff.Dual}
     _restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
 end
 
+function get_tmp(dc::DiffCache, u::Type{T}) where {T <: ForwardDiff.Dual}
+    nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+    if nelem > length(dc.dual_du)
+        enlargediffcache!(dc, nelem)
+    end
+    _restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+end
+
 function get_tmp(dc::DiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual}
     nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
     if nelem > length(dc.dual_du)