diff src/LazyTensors/lazy_tensor_operations.jl @ 919:b41180efb6c2 performance/get_region_type_inference

Start refactor to improve type stability for apply(op,::TensorApplication,...)
author Jonatan Werpers <jonatan@werpers.com>
date Mon, 21 Feb 2022 10:33:58 +0100
parents 76e5682d0e52
children 8b0ff2fddc32 f1c2a4fa0ee1
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Tue Feb 15 15:14:28 2022 +0100
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Mon Feb 21 10:33:58 2022 +0100
@@ -430,3 +430,28 @@
     print(io, "SizeMismatch: ")
     print(io, "domain size $(domain_size(err.tm)) of TensorMapping not matching size $(err.sz)")
 end
+
+function apply_with_region(op, v, boundary_width::Integer, dim_size::Integer, i)
+    if 0 < i <= boundary_width
+        return LazyTensors.apply(op,v,Index(i,Lower))
+    elseif boundary_width < i <= dim_size-boundary_width
+        return LazyTensors.apply(op,v,Index(i,Interior))
+    elseif dim_size-boundary_width < i <= dim_size
+        return LazyTensors.apply(op,v,Index(i,Upper))
+    else
+        error("Bounds error") # TODO: Make this more standard
+    end
+end
+# TBD: Can these methods be merge by having a function as an arguement instead?
+# TODO: Add inference test that show where things break and how this rewrite fixes it.
+function apply_transpose_with_region(op, v, boundary_width::Integer, dim_size::Integer, i)
+    if 0 < i <= boundary_width
+        return LazyTensors.apply_transpose(op,v,Index(i,Lower))
+    elseif boundary_width < i <= dim_size-boundary_width
+        return LazyTensors.apply_transpose(op,v,Index(i,Interior))
+    elseif dim_size-boundary_width < i <= dim_size
+        return LazyTensors.apply_transpose(op,v,Index(i,Upper))
+    else
+        error("Bounds error") # TODO: Make this more standard
+    end
+end