Fix Backend Compatibility

[alyst]

Fixes #13 (matrix multiplication tests succeed on ReverseDiff.jl, but / tests fail because it's not supported yet (IIUC))

Closes #15 (switches to AbstractArray)

Moves sparse matrix tests into extension (requires MKLSparse.jl).

Tests now include checks for type stability of the test functions.

[devmotion]

Can you explain why you switched the constant tuples to arrays?

[alyst]

For consistency with matrix2matrix functions, which are extended in DiffTestsMKLSparseExt. I can revert num2num to be tuples again.

But out of curiosity - why they were tuples in the first place? What kind of optimisation does it trigger?

[devmotion]

I assume because using tuples reduces allocations.

I don't think it matters much but I'd suggest reverting these changes in the PR because 1) they seem unrelated and 2) it is not clear that they improve anything (on the contrary, they increase allocations).

[alyst]

I can revert it for the sets that are not modified in the extension. BTW, how much would it reduce memory overhead (considering that in return Julia would have to construct a tuple type)?

[devmotion]

🤷

[devmotion]

I'm a fan of keeping changes minimal and avoiding unrelated changes. So I'd suggest reverting these changes, regardless of whether they can/should be changed or not. For consistency, I think it would make sense to just use a tuple in the extension as well - and to open an issue or separate PR if you would like to change these tuples to arrays.

[alyst]

It is possible to replace a tuple with the extended one (although I would then have to drop const from the declaration); the code would look a little bit more convoluted.
But if you mean that the extension should define its own set of sparse functions to test -- that would require adding support for that set of functions in all AD backend tests suites.

I am also a fan of minimal changes, but there are also consistency and maintainability considerations. For this PR I thought it would be strange to change some function sets into arrays and keep others as tuples, or rewrite tuples from the extension. But if you, as a maintainer, prefer minimal changes, I can minimize them.

[devmotion]

Oh, I missed that the extension modifies arrays in the package. That does not seem OK to me, I don't think extensions are supposed to do that. For instance, this means that the resulting state depends on the order in which (possibly additional) extensions are loaded etc.

So in case you want to add an extension, I think a different design is needed anyway.

And yes, I very much would prefer more minimal changes.

[alyst]

That does not seem OK to me, I don't think extensions are supposed to do that. For instance, this means that the resulting state depends on the order in which (possibly additional) extensions are loaded etc.

What is the specific scenario that you think would be problematic?
The extensions would not remove functions from the lists, and they would be checking if the function is already in the list.
As for the order -- I don't think the order of functions in the list has any noticeable effects.
One thing to keep in mind is Revise.jl, but at the moment DiffTests.jl declares function lists as const, so Revise.jl cannot update current lists on the fly -- there is no regression here.

Also vector-based approach is used (see below) to dynamically populate the list based on Julia version (e.g. because some linear algebra routines are not available in the earlier versions). DiffTests.jl may drop support for Julia 1.0, but in the long run it might be beneficial to have some straightforward mechanism for supporting multiple Julia versions.

And let me refer to the original issue that I am trying to solve -- I've created JuliaDiff/ForwardDiff.jl#589 to address IMO the very relevant practical problem -- improving autodiff efficiency for linear models and multivariate normal distributions with dense or sparse matrices. So the test functions that I am adding here via the extension should be relevant for all autodiff backends. Vanilla SparseArrays.jl does not, at the moment, support ldiv!() for triangular sparse matrices, hence MKLSparse.jl weak dependency.

[devmotion]

Extensions are just not designed for such a use case - their sole purpose is to extend existing functions.

A problematic scenario here would be another extension that also adds functions to the same array. Then the order of the functions in the array after loading the extensions can be completely arbitrary, and even worse if precompilation and loading of the extensions happens with multiple tasks in parallel two push!es might interfere.

[devmotion]

BTW this should not be done: Extensions should only import the package itself and the weak dependency but not even standard libraries.

[alyst]

Why? SparseArrays and LinearAlgebra are the direct dependencies of DiffTests itself and MKLSparse. The latter provides more specific definitions of matrix multiplication and division for sparse matrices that are utilized by the test functions.

[devmotion]

You should import them indirectly via DiffTests and MKLSparse (see eg JuliaStats/LogExpFunctions.jl#63).

[alyst]

Thank you for the info. I've updated how the import is being done. However, the changes that the extensions makes to the function lists are not visible from the user session (I guess it's different world age of DiffTests.jl or something like that).

[devmotion]

To fix #17 (comment) properly, I had the following idea: Could we add traits to the test functions (eg., istestfunc, isinplace, inputdim, outputdim or whatever is appropriate) plus some convenience functions number_to_number_funcs etc. that basically just filter every function (or only inside of DiffTests?) according to the desired criteria? Then we could avoid pushing to fixed containers completely and would only extend the traits in the extension.

[alyst]

Could we add traits to the test functions (eg., istestfunc, isinplace, inputdim, outputdim or whatever is appropriate) plus some convenience functions number_to_number_funcs etc

I had a similar idea. But then, IIUC, to collect those functions one would have to use methods() -- this works fine, but the result is MethodList, and the elements are of type Method, not Function. Method is not callable, and at the first glance it looked like conversion would require some quirky low-level manipulations.

I figured out the problem I had with extension not updating the list of functions. It turned out that with each import/using DiffTests instruction, the DiffTests module was re-evaluated, i.e. all const assignments were reset. To fix that, I have moved assignments to __init__() (both in DiffTests and the extensions), which is called only once. Adding functions to the list is done with DiffTests.register(registry_id, functions) method, i.e. DiffTests.register(:VECTOR_TO_VECTOR_FUNCS, [diag_lmul, utriag_lmul]). The method checks that no duplicates are being introduced, and I can add critical section to ensure that parallel threads would not access it simultaneously. Under the hood, register() uses setproperty!() (wrapper for setglobal!() on 1.9, or setfield!() in earlier versions) to declare the corresponding global variable in DiffTests module, so no changes are required downstream.

[alyst]

@devmotion I think the issue with updating the function lists by the extensions is solved. Do you have any further feedback on the PR? The functions functions that are being enabled are important for multivariate distributions.

[devmotion]

Based on the previous comment, my impression is that there are still problems with e.g. simultaneous register calls. I'm generally skeptical regarding such global lists based on previous experiences with e.g. DiffRules. Lastly, I'm a bit concerned adding a (weak) dependency that is broken on MacBooks (precompilation errors):

(jl_L4HrUa) pkg> precompile MKLSparse
Precompiling MKLSparse
  ✗ MKLSparse
  0 dependencies successfully precompiled in 1 seconds. 4 already precompiled.

ERROR: The following 1 direct dependency failed to precompile:

MKLSparse [0c723cd3-b8cd-5d40-b370-ba682dde9aae]

Failed to precompile MKLSparse [0c723cd3-b8cd-5d40-b370-ba682dde9aae] to "/Users/david/.julia/compiled/v1.9/MKLSparse/jl_aUHjJZ".
ERROR: LoadError: UndefVarError: `libmkl_rt` not defined
Stacktrace:
 [1] include(mod::Module, _path::String)
   @ Base ./Base.jl:457
 [2] include(x::String)
   @ MKLSparse ~/.julia/packages/MKLSparse/4LlPk/src/MKLSparse.jl:1
 [3] top-level scope
   @ ~/.julia/packages/MKLSparse/4LlPk/src/MKLSparse.jl:9
 [4] include
   @ ./Base.jl:457 [inlined]
 [5] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
   @ Base ./loading.jl:2049
 [6] top-level scope
   @ stdin:3
in expression starting at /Users/david/.julia/packages/MKLSparse/4LlPk/src/BLAS/BLAS.jl:1
in expression starting at /Users/david/.julia/packages/MKLSparse/4LlPk/src/MKLSparse.jl:1
in expression starting at stdin:3

[alyst]

@devmotion Thanks for the feedback!

I'm generally skeptical regarding such global lists based on previous experiences with JuliaDiff/DiffRules.jl#90.

I agree that for highly extendable functionality, like DiffRules.jl, the current mechanism might be not flexible enough.
But our scope here is much more humble -- it is the package-internal mechanism to update lists of functions from its extensions (the register() method is not exported).
The implementation already updates the function lists from the extension's __init__() (the major limitation listed in the issue you referenced).
Basically, the only way the person can influence the lists of tested functions is via the project dependencies (which would activate the corresponding extensions).

my impression is that there are still problems with e.g. simultaneous register calls.

I can add locks to ensure exclusive access if that is what you mean.

Lastly, I'm a bit concerned adding a (weak) dependency that is broken on MacBooks (precompilation errors)

AFAIK there are issues with Mac version of MKLSparse with some of the linear algebra functions (see e.g. JuliaSparse/MKLSparse.jl#37), but the HEAD passes the tests.
The last release is 3 year-old, though (I am trying to fix that too). Which MKLSparse.jl version have you tried?
It's a bit hard to test across the architectures with GitHub CI right now, since it requires new releases of both MKLSparse.jl and DiffTests.jl.

[devmotion]

MKL is not available for ARM, everyone with Apple silicon will run into the same issue (see, e.g., JuliaMath/FFTW.jl#274).

[devmotion]

Which makes me wonder: Maybe we should just not use MKL, and hence maybe we should just not add an extension. And thereby we could get rid of the design issues as well.

[alyst]

MKL is not available for ARM, everyone with Apple silicon will run into the same issue (see, e.g., JuliaMath/FFTW.jl#274).
Which makes me wonder: Maybe we should just not use MKL, and hence maybe we should just not add an extension. And thereby we could get rid of the design issues as well.

1. I would too prefer to avoid MKLSparse.jl dependency, but at the moment one runs into problems without it when dealing with certain models that have sparse covariation matrices.
2. I guess one of the goals of JuliaDiff ecosystem is to be extendable (or at least to support a wide range of applications). My experience is that right now it is not so easy.
   The PR is my attempt to improve it a little bit by providing a generic mechanism to update the list of tested functions based on the project dependencies.
   It is not flexible, but it is minimally disruptive to the current behaviour and package structure -- if one does not add these packages to the [test] project dependencies, nothing changes.