Overloading, vector inputs

[aml5600]

Hello! I am (attempting) to solve a large NLP through Optimization.jl. The problem is quite sparse, so using AutoSparse for AD has been quite helpful and I see that you guys are to thank for that, so thank you!

A part of my NLP is doing inference on an already trained Lux model. When the overall function is being traced, I get a variety of warnings about fallback matmul because of mismatched types (one being floats, the other being the tracers). Previously (I forgot I did this...) I overloaded my wrapper of the Lux model on the tracer types and essentially did ones(output_size, input_size) * x, so that the calling function gave full incidence on the input. I am not too worried about getting the exact sparsity through the entire model; the input/output sizes are small relative to some of the internal dimensions.

I am now trying to do this correctly using the overloads that the package provides. The examples seem to only touch on scalar inputs. Is there anything special/different that must be done for (say) a vector of tracer types as the input? From above, x was overloaded on an AbstractVector of the tracer type. Just want to ensure this is the correct way forward.

Thanks again,
Andrew

[adrhill]

Hi Andrew, thanks for the feedback!

I am now trying to do this correctly using the overloads that the package provides. The examples seem to only touch on scalar inputs.

Indeed, our focus is on scalar overloads. However, for global sparsity detection (TracerSparsityDetector) we also provide some array-level overloads. These two files should provide you with examples:

• Overloads on common LinearAlgebra.jl functions: https://github.com/adrhill/SparseConnectivityTracer.jl/blob/e6d357b5e0b23d72ab3430e359831eafeb005d73/src/overloads/arrays.jl
• Our DataInterpolations.jl extension: https://github.com/adrhill/SparseConnectivityTracer.jl/blob/main/ext/SparseConnectivityTracerDataInterpolationsExt.jl

A pattern that you will see repeated is that we take a union of the input tracers and return a FillArray of the union of tracers. To understand what these unions do, I recommend reading our "How it works guide", especially the section on index sets.
We are able to do this because we allow ourselves to return overly conservative patterns. (Contributions to return less convervative patterns are always welcome.)
For this purpose, we provide two helper functions:

SparseConnectivityTracer.jl/src/overloads/utils.jl

Lines 5 to 48 in e6d357b

	"""
	first_order_or(tracers)
	Compute the most conservative elementwise OR of tracer sparsity patterns,
	excluding second-order interactions of `HessianTracer`.
	This is functionally equivalent to:
	```julia
	reduce(+, tracers)
	```
	"""
	function first_order_or(ts::AbstractArray{T}) where {T<:AbstractTracer}
	# TODO: improve performance
	return reduce(first_order_or, ts; init=myempty(T))
	end
	function first_order_or(a::T, b::T) where {T<:GradientTracer}
	return gradient_tracer_2_to_1(a, b, false, false)
	end
	function first_order_or(a::T, b::T) where {T<:HessianTracer}
	return hessian_tracer_2_to_1(a, b, false, true, false, true, true)
	end
	"""
	second_order_or(tracers)
	Compute the most conservative elementwise OR of tracer sparsity patterns,
	including second-order interactions to update the `hessian` field of `HessianTracer`.
	This is functionally equivalent to:
	```julia
	reduce(^, tracers)
	```
	"""
	function second_order_or(ts::AbstractArray{T}) where {T<:AbstractTracer}
	# TODO: improve performance
	return reduce(second_order_or, ts; init=myempty(T))
	end
	function second_order_or(a::T, b::T) where {T<:GradientTracer}
	return gradient_tracer_2_to_1(a, b, false, false)
	end
	function second_order_or(a::T, b::T) where {T<:HessianTracer}
	return hessian_tracer_2_to_1(a, b, false, false, false, false, false)
	end

You can see these helper functions and FillArrays being used in the following example:

SparseConnectivityTracer.jl/src/overloads/arrays.jl

Lines 101 to 106 in e6d357b

	function Base.exp(A::AbstractMatrix{T}) where {T<:AbstractTracer}
	LinearAlgebra.checksquare(A)
	n = size(A, 1)
	t = second_order_or(A)
	return Fill(t, n, n)
	end

When the overall function is being traced, I get a variety of warnings about fallback matmul because of mismatched types (one being floats, the other being the tracers).

Could you open an issue separate issue for this? I think it would be desirable for us to support Lux.jl out of the box.

[aml5600]

Incredible response, thanks! the examples from utils.jl and arrays.jl answered my question. I will indeed open an issue for a Lux extension.