Add Tensor wrapper for jacobians in DualArrays

Project.toml

@@ -1,7 +1,7 @@
 name = "DualArrays"
 uuid = "429e4e16-f749-45f3-beec-30742fae38ce"
-authors = ["Sheehan Olver <solver@mac.com>"]
 version = "0.2.0"
+authors = ["Sheehan Olver <solver@mac.com>"]
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -10,13 +10,15 @@ DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
 [compat]
 DiffRules = "1.15.1"
 ForwardDiff = "1.2.2"
 Lux = "1.29.1"
 Plots = "1.41.1"
 SparseArrays = "1.10"
+TensorOperations = "5.5.2"
 
 [extras]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"

src/DualArrays.jl

@@ -6,16 +6,18 @@ A Julia package for efficient automatic differentiation using dual numbers and d
 This package provides:
 - `Dual`: A dual number type for storing values and their derivatives
 - `DualVector`: A vector of dual numbers represented with a Jacobian matrix
-
+- `DualMatrix`: A matrix of dual numbers represented with a Jacobian tensor`
 Differentiation rules are mostly provided by ChainRules.jl.
 """
 module DualArrays
 
-export DualVector, Dual, jacobian
+export DualVector, DualMatrix, Dual, jacobian
+
+import Base: +, -, ==, getindex, size, axes, broadcasted, show, sum, vcat, convert, *, isapprox, \, eltype, transpose, permutedims
 
-import Base: +, -, ==, getindex, size, axes, broadcasted, show, sum, vcat, convert, *, isapprox, \
+using LinearAlgebra, ArrayLayouts, FillArrays, DiffRules, TensorOperations
 
-using LinearAlgebra, ArrayLayouts, FillArrays, DiffRules
+import FillArrays: elconvert
 
 include("types.jl")
 include("indexing.jl")

src/arithmetic.jl

@@ -96,4 +96,66 @@ LinearAlgebra.dot(x::DualVector, y::AbstractVector) = Dual(dot(x.value, y), tran
 LinearAlgebra.dot(x::AbstractVector, y::DualVector) = Dual(dot(x, y.value), transpose(y.jacobian) * x)
 
 # solve
-\(x::AbstractMatrix, y::DualVector) = DualVector(x \ y.value, x \ y.jacobian)
+\(x::AbstractMatrix, y::DualVector) = DualVector(x \ y.value, ArrayOperator{1}(x \ y.jacobian.data))
+
+"""
+-----------------------
+Multiplication with ArrayOperators
+-----------------------
+
+* for ArrayOperators generalises from matrix/vector
+multiplication and uses tensor contraction provided by
+TensorOperations.jl. Let an (A, B) ArrayOperator denote
+an ArrayOperator with output dimension A and input dimension B
+(i.e an ArrayOperator{A+B, T, A, B} for some type T). We only
+allow multiplication between an (A, B) ArrayOperator and a (B, C)
+ArrayOperator, with the result being an (A, C) ArrayOperator.
+This generalises from:
+
+-An inner product: a row vector * a column vector -> a scalar:
+    (0, 1) * (1, 0) -> (0, 0)
+
+-An outer product: a column vector * a row vector -> a matrix:
+    (1, 0) * (0, 1) -> (1, 1)
+
+- matrix/vector multiplication: a matrix * a column vector -> a column vector:
+    (1, 1) * (1, 0) -> (1, 0)
+
+We can treat multiplication between ArrayOperators and regular arrays as multiplication
+between two ArrayOperators: given an (A, B) ArrayOperator and an L-array, we can fix it as
+having input dimension B and infer C = L - B. We use this contraction rule to multiply the two.
+"""
+
+# Helper function to perform tensor contraction between x and y
+function _contract(x, y, A, B, C)
+    x_idx = (ntuple(i -> i, A), ntuple(i -> i + A, B))
+    y_idx = (ntuple(i -> i, B), ntuple(i -> i + B, C))
+    ret_idx = (ntuple(i -> i, A), ntuple(i -> i + A, C))
+
+    return TensorOperations.tensorcontract(x, x_idx, false, y, y_idx, false, ret_idx, 1)
+end
+
+function *(x::ArrayOperator{A, B}, y::ArrayOperator{B, C}) where {A, B, C}
+    return ArrayOperator{A}(_contract(x.data, y.data, A, B, C))
+end
+
+function *(x::ArrayOperator{A, B}, y::AbstractArray{<:Any, L}) where {A, B, L}
+    C = L - B
+    return ArrayOperator{A}(_contract(x.data, y, A, B, C))
+end
+
+function *(x::AbstractArray{<:Any, L}, y::ArrayOperator{B, C}) where {L, B, C}
+    A = L - B
+    return ArrayOperator{A}(_contract(x, y.data, A, B, C))
+end
+
+# Extra required arithmetic for ArrayOperators
+Base.:+(x::ArrayOperator, y::ArrayOperator) = x .+ y
+Base.:-(x::ArrayOperator, y::ArrayOperator) = x .- y
+Base.:-(x::ArrayOperator) = (-).(x)
+
+Base.:*(a::Number, t::ArrayOperator{N, M, T, L}) where {N, M, T, L} =
+    ArrayOperator{N, M, promote_type(typeof(a), T), L}(a .* t.data)
+Base.:*(t::ArrayOperator, a::Number) = a * t
+
+

src/indexing.jl

@@ -6,28 +6,62 @@ sparse_getindex(a...) = layout_getindex(a...)
 sparse_getindex(D::Diagonal, k::Integer, ::Colon) = OneElement(D.diag[k], k, size(D, 2))
 sparse_getindex(D::Diagonal, ::Colon, j::Integer) = OneElement(D.diag[j], j, size(D, 1))
 
+
+# We need a system of indexing that takes two tuples
+# of length N and M. We must then return a ArrayOperator whose new input and output dimensions
+# are inferred from how many of the arguments in each respective tuple are integers.
+
+# For the purposes of DualArrays.jl, we only need to index the input dimensions
+#
+# Example:
+#
+# Consider a DualVector with standard matrix Jacobian (ArrayOperator{2, T, 1, 1}):
+#
+# d = a + Bϵ
+#
+# If we index d[1], we have:
+#
+# d[1] = a[1] + B[1, :]ϵ
+#
+# Where B[1, :] indexes the input dimension of B
+Idx = Union{Int, Colon, AbstractUnitRange}
+count_ints(t::Tuple) = count(x -> x isa Int, t)
+
+function getindex(t::ArrayOperator{N, M}, i::NTuple{N, Idx}, j::NTuple{M, Idx}) where {N, M}
+    # new value of M is inferred from the ArrayOperator constructor and the order of
+    # indexing the underlying array.
+    newN = N - count_ints(i)
+    ArrayOperator{newN}(sparse_getindex(t.data, i..., j...))
+end
+
+function getindex(t::ArrayOperator{N, M}, i::NTuple{N, Idx}, ::Colon) where {N, M}
+    # new value of M is inferred from the ArrayOperator constructor and the order of
+    # indexing the underlying array.
+    newN = N - count_ints(i)
+    ArrayOperator{newN}(sparse_getindex(t.data, i..., ntuple(_ -> Colon(), M)...))
+end
+# Integer indexing always returns a scalar.
+getindex(t::ArrayOperator, i::Vararg{Int}) = getindex(t.data, i...)
+# Indexing with only slices/ranges returns a similar tensor
+getindex(t::ArrayOperator{N, M}, i::Vararg{Union{Colon, UnitRange}}) where {N, M} = ArrayOperator{N}(sparse_getindex(t.data, i...))
+
 """
-Extract a single Dual number from a DualVector at position y.
+Extract a single Dual number from a DualArray.
 """
-function Base.getindex(x::DualVector, y::Int)
-    Dual(x.value[y], sparse_getindex(x.jacobian, y, :))
+function getindex(x::DualArray, args::Vararg{Int})
+    Dual(x.value[args...], x.jacobian[args, :])
 end
 
 """
 Extract a sub-DualVector from a DualVector using a range.
 """
-function Base.getindex(x::DualVector, y::UnitRange)
+function getindex(x::DualVector, y::UnitRange)
     newval = x.value[y]
-    newjac = sparse_getindex(x.jacobian, y, :)
+    newjac = x.jacobian[y, :]
     DualVector(newval, newjac)
 end
 
-"""
-Return the size of the DualVector (length of the value vector).
-"""
-Base.size(x::DualVector) = (length(x.value),)
-
-"""
-Return the axes of the DualVector.
-"""
-Base.axes(x::DualVector) = axes(x.value)
+# DualArray array interface
+for op in (:size, :axes)
+    @eval $op(a::DualArray) = $op(a.value)
+end