Should `_trunc` methods accept a `Vector` or `Ref` argument for the truncation error?

[kshyatt]

This would make writing Enzyme rules much easier as it would allow us to avoid mixed activity. It would also help avoid GPU synchronization by storing the truncation error in a GPU-side vector which can be read if necessary. Would it make sense to add another interface for the svd_trunc!, eig_trunc!, and friends like:

svd_trunc!(A, USVh, epsilon::AbstractVector{T}; alg) where {T<:Real}, where epsilon is length 1 and overwritten?

[kshyatt]

cc @lkdvos and @Jutho

[lkdvos]

Apologies for taking a while to get to this.

While I can see this being convenient from an implementation point of view, I do feel like it is a bit unexpected from a user point of view to return a container for something that really is a scalar.
Does Enzyme have an "overlay" mode where we can swap out a function by a differentiable one instead, so we only have to do this when AD is active?
Alternatively, would you suggest we do this always, or simply only for the internal implementation, and then unwrap this afterwards?

[Jutho]

I haven't thought strongly about this, but another benefit might be that epsilon could then be a gpu array and so there would not be any actual data transfer between GPU and CPU within the scope of this call (I think).

I only think about this because I was reminded about the KrylovKit issue on GPUs, where at some point we will also need to have an inner method that stores the variable in a reference or array, in order to avoid having to communicate back and forth between GPU and CPU for the Gram-Schmidt process. I am not sure something similar applies here, so maybe this is off topic.

However, I think I would prefer epsilon to become part of the output tuple if we adopt this interface. The only problem then is that it is probably more breaking, as otherwise we can just add this method with a default argument for epsilon. However, also with the tuple this can be done I think.

[kshyatt]

Does Enzyme have an "overlay" mode where we can swap out a function by a differentiable one instead, so we only have to do this when AD is active?

I don't believe they have overlay mode

Alternatively, would you suggest we do this always, or simply only for the internal implementation, and then unwrap this afterwards?

I think doing it internally could work. I believe that's the suggestion Enzyme itself makes in https://enzymead.github.io/Enzyme.jl/dev/faq/#Mixed-activity

[lkdvos]

Copying @Jutho's suggestion here:

We have a new method svd_trunc!(A, (U, S, Vᴴ, ϵ), alg) with a length-four tuple as second argument, where ϵ is a vector of the same type of S with length either one or zero. If length is zero, we do not compute the error, if length is one, we do. The return value of this method is (U, S, Vᴴ, ϵ).

Then, the current method svd_trunc!(A, (U, S, Vᴴ), alg) is changed into svd_trunc!(A, (U, S, Vᴴ), alg; compute_error = true)

if compute_error, we allocate a length-one vector ϵ and call svd_trunc!(A, (U, S, Vᴴ, ϵ), alg) which returns the output (U, S, Vᴴ, ϵ). We capture that output and return U, S, Vᴴ, ϵ[] instead.

if !compute_error, we allocate a length-zero vector ϵ, and call
svd_trunc!(A, (U, S, Vᴴ, ϵ), alg) which returns the (U, S, Vᴴ, ϵ). We capture that output and return something like U, S, Vᴴ, NaN.

---

addendum from myself:

We should be careful to ensure that also downstream packages that implement their own alg types still work as intended, and it could be good to explicitly check that (i.e. by running downstream tests, as suggested by @kshyatt ).