add type annotation to _evaluate function

rittermarc · rittermarc · commit d4ef4e30b6a6 · 2024-03-21T14:55:33.000-04:00
diff --git a/src/tensortrain.jl b/src/tensortrain.jl
@@ -88,7 +88,9 @@ function to_tensors(obj::TensorTrainFit{ValueType}, x::Vector{ValueType}) where
     ]
 end
 
-_evaluate(tt, indexset) = only(prod(T[:, i, :] for (T, i) in zip(tt, indexset)))
+function _evaluate(tt::Vector{Array{V, 3}}, indexset) where {V}
+    only(prod(T[:, i, :] for (T, i) in zip(tt, indexset)))
+end
 
 function (obj::TensorTrainFit{ValueType})(x::Vector{ValueType}) where {ValueType}
     tensors = to_tensors(obj, x)

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,9 @@ function to_tensors(obj::TensorTrainFit{ValueType}, x::Vector{ValueType}) where`
`88`	`88`	`]`
`89`	`89`	`end`
`90`	`90`
`91`		`-_evaluate(tt, indexset) = only(prod(T[:, i, :] for (T, i) in zip(tt, indexset)))`
	`91`	`+function _evaluate(tt::Vector{Array{V, 3}}, indexset) where {V}`
	`92`	`+ only(prod(T[:, i, :] for (T, i) in zip(tt, indexset)))`
	`93`	`+end`
`92`	`94`
`93`	`95`	`function (obj::TensorTrainFit{ValueType})(x::Vector{ValueType}) where {ValueType}`
`94`	`96`	`tensors = to_tensors(obj, x)`