Skip to content

Add rules and tests for kron #741

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 16 commits into
base: main
Choose a base branch
from
Open

Add rules and tests for kron #741

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
5bb9766
Add rules and tests
Feb 6, 2024
f0902e3
Add tests for `rrule`
Feb 6, 2024
7c53f4b
Add rules and try to cover complex case
Feb 6, 2024
c1226eb
Restrict types of arguments
simsurace Feb 6, 2024
236daf1
Write rules functionally and fix them
Feb 6, 2024
b2d4f4a
Add `unthunk` and `@thunk`
Feb 6, 2024
8b94cfc
Change dimensions to make them recognizable
Feb 6, 2024
b71b8ef
Further simplify rules
Feb 6, 2024
2ad5473
Only define rules for Julia 1.9 onwards
Feb 6, 2024
fde509e
Add projections
Feb 6, 2024
4386143
Add projections and remove redundant `conj` calls
Feb 6, 2024
1b97828
Fix type instability
Feb 6, 2024
5b74071
Run tests only above Julia 1.9
simsurace Feb 6, 2024
72060d0
Fix typo
Feb 6, 2024
6650adc
Enable all tests
Feb 6, 2024
f104172
Improve version bound
simsurace Feb 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 61 additions & 0 deletions src/rulesets/LinearAlgebra/dense.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -394,3 +394,64 @@ function rrule(
end
return Ω, lyap_pullback
end

#####
##### `kron`
#####

@static if VERSION ≥ v"1.9.0-DEV.1267"
function frule((_, Δx, Δy), ::typeof(kron), x::AbstractVecOrMat{<:Number}, y::AbstractVecOrMat{<:Number})
return kron(x, y), kron(Δx, y) + kron(x, Δy)
end

function rrule(::typeof(kron), x::AbstractVector{<:Number}, y::AbstractVector{<:Number})
project_x = ProjectTo(x)
project_y = ProjectTo(y)
function kron_pullback(z̄)
dz = reshape(unthunk(z̄), length(y), length(x))
x̄ = @thunk(project_x(conj.(dz' * y)))
ȳ = @thunk(project_y(dz * conj.(x)))
return NoTangent(), x̄, ȳ
end
return kron(x, y), kron_pullback
end

function rrule(::typeof(kron), x::AbstractMatrix{<:Number}, y::AbstractVector{<:Number})
project_x = ProjectTo(x)
project_y = ProjectTo(y)
function kron_pullback(z̄)
dz = reshape(unthunk(z̄), length(y), size(x)...)
x̄ = @thunk(project_x(dot.(Ref(y), eachslice(dz; dims = (2, 3)))))
ȳ = @thunk(project_y(dot.(Ref(x), eachslice(dz; dims = 1))))
return NoTangent(), x̄, ȳ
end
return kron(x, y), kron_pullback
end

function rrule(::typeof(kron), x::AbstractVector{<:Number}, y::AbstractMatrix{<:Number})
project_x = ProjectTo(x)
project_y = ProjectTo(y)
function kron_pullback(z̄)
dz = reshape(unthunk(z̄), size(y, 1), length(x), size(y, 2))
x̄ = @thunk(project_x(dot.(Ref(y), eachslice(dz; dims = 2))))
ȳ = @thunk(project_y(dot.(Ref(x), eachslice(dz; dims = (1, 3)))))
return NoTangent(), x̄, ȳ
end
return kron(x, y), kron_pullback
end

function rrule(::typeof(kron), x::AbstractMatrix{<:Number}, y::AbstractMatrix{<:Number})
project_x = ProjectTo(x)
project_y = ProjectTo(y)
function kron_pullback(z̄)
dz = reshape(unthunk(z̄), size(y, 1), size(x, 1), size(y, 2), size(x, 2))
x̄ = @thunk(project_x(_dot_collect.(Ref(y), eachslice(dz; dims = (2, 4)))))
ȳ = @thunk(project_y(_dot_collect.(Ref(x), eachslice(dz; dims = (1, 3)))))
Comment on lines +446 to +449
Copy link
Member

@mcabbott mcabbott Oct 3, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was wondering if you have to make slices, given that kron is just reshape and .*. So here's an attempt to do without:

using ChainRulesCore

function pr_rule(x::AbstractMatrix{<:Number}, y::AbstractMatrix{<:Number})  # from https://github.com/JuliaDiff/ChainRules.jl/pull/741
    project_x = ProjectTo(x)
    project_y = ProjectTo(y)
    function kron_pullback(z̄)
        dz = reshape(unthunk(z̄), size(y, 1), size(x, 1), size(y, 2), size(x, 2))
        x̄ = @thunk(project_x(dot.(Ref(y), eachslice(dz; dims = (2, 4)))))
        ȳ = @thunk(project_y(dot.(Ref(x), eachslice(dz; dims = (1, 3)))))
        return NoTangent(), x̄, ȳ
    end
end

# using TensorCast
# mykron(x,y) = @cast z[(a,b), (c,d)] := x[b,d] * y[a,c]
# @pretty @cast z[(a,b), (c,d)] := x[b,d] * y[a,c]

function shape_rule(x::AbstractMatrix, y::AbstractMatrix)
    function back(dz)
        x4 = reshape(x, 1, size(x,1), 1, size(x,2))
        y4 = reshape(y, size(y,1), 1, size(y,2), 1)
        dz4 = reshape(unthunk(dz), size(y,1), size(x,1), size(y,2), size(x,2))
        dx = @thunk ProjectTo(x)(reshape(sum(dz4 .* y4, dims=(1,3)), size(x)))  # might be missing conj
        dy = @thunk ProjectTo(y)(reshape(sum(dz4 .* x4, dims=(2,4)), size(y)))
        0, dx, dy
      end
end

let x = rand(10,20), y = rand(30,10)
  b1 = pr_rule(x, y)
  b2 = shape_rule(x, y)
  z = kron(x,y)
  _, dx1, _ = @btime map(unthunk, $b1($z))
  _, dx2, _ = @btime map(unthunk, $b2($z))
  dx1 ≈ dx2
end
  # min 181.458 μs, mean 185.668 μs (4 allocations, 4.39 KiB)
  # min 80.583 μs, mean 169.305 μs (32 allocations, 943.05 KiB)
# true

It's a pity to allocate these big arrays dz4 .* y4 but still seems quicker. Possibly we could use lazy broadcasting to avoid that:

bc = Broadcast.instantiate(Broadcast.broadcasted(*, [1 2 3], [4, 5]));
sum(bc)  # OK
sum(bc; dims=1) # ERROR: MethodError: no method matching reducedim_init(::typeof(identity), ::typeof(Base.add_sum), ::Base.Broadcast.Broadcasted{…}, ::Int64)
sum!([0 0 0], bc)  # ERROR: MethodError: no method matching sum!(::Matrix{Int64}, ::Base.Broadcast.Broadcasted
sum(bc; dims=1, init=0.0)  # OK, not sure if it's fast or not

On StaticArrays (mentioned above) both at present make a SizedMatrix, which I think is ProjectTo's attempt to fix things up. Surely this reshaping could be done in a static-friendly way but IDK exactly how.

julia> let x = @SMatrix(rand(5,5)), y = @SMatrix(rand(5,5))
         b1 = pr_rule(x, y)
         b2 = shape_rule(x, y)
         z = kron(x,y)
         _, dx1, _ = @btime map(unthunk, $b1($z))
         _, dx2, _ = @btime map(unthunk, $b2($z))
         dx1 ≈ dx2
       end
  min 2.458 μs, mean 2.558 μs (2 allocations, 512 bytes)
  min 4.006 μs, mean 5.198 μs (22 allocations, 11.38 KiB)
true

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this result scale to larger arrays?

Copy link
Member

@mcabbott mcabbott Oct 5, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Result meaning speed difference? It will vary with size & machine. On very small arrays reshaping is faster slower! (Like 3x3 I meant.)

Issues with StaticArrays will be similar at all sizes.

I think broadcasting over slices will work badly on CuArrays, and tend to make Arrays. But right now neither idea seems to work, not sure why

julia> using Metal

julia> bk = pr_rule(MtlArray(rand(Float32, 3,3)), MtlArray(rand(Float32, 3,3)));

julia> bk(MtlArray(rand(Float32, 9,9)))[2] |> unthunk
ERROR: GPU compilation of MethodInstance for (::GPUArrays.var"#broadcast_kernel#26")(::Metal.mtlKernelContext, ::MtlDeviceMatrix{…}, ::Base.Broadcast.Broadcasted{…}, ::Int64) failed
KernelError: passing and using non-bitstype argument

julia> bk2 = shape_rule(MtlArray(rand(Float32, 3,3)), MtlArray(rand(Float32, 3,3)));

julia> bk2(MtlArray(rand(Float32, 9,9)))[2] |> unthunk
ERROR: could not load symbol "LLVMExtraAddPropagateJuliaAddrspaces":
dlsym(RTLD_DEFAULT, LLVMExtraAddPropagateJuliaAddrspaces): symbol not found

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the reshape version is not strictly better than the current one, especially for large arrays, I would propose to keep the current version and put further optimizations in a separate PR.

Copy link
Member

@mcabbott mcabbott Oct 5, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A bit curious at what sizes it's slower for you?

But mainly I think the issue is less about the race than that simple solid-array operations have a better chance of behaving well with StaticArrays, and CuArrays. I haven't taken another pass to see if the first draft can be improved on.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't benchmarked anything myself yet. I will give it a go later.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, results seem to be mixed. For larger sizes the allocations are taking their price:

let x = rand(100,200), y = rand(300,100)
    b1 = pr_rule(x, y)
    b2 = shape_rule(x, y)
    z = kron(x,y)
    _, dx1, _ = @btime map(unthunk, $b1($z))
    _, dx2, _ = @btime map(unthunk, $b2($z))
    dx1 ≈ dx2
end
#   3.376 s (6 allocations: 390.84 KiB)
#   3.797 s (34 allocations: 8.94 GiB)
# true

I would suggest staying with the current implementation.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One way to ensure any implementation isn't excluding all GPU array types would be to toss a @gpu in front of the new tests, no?

return NoTangent(), x̄, ȳ
end
return kron(x, y), kron_pullback
end

_dot_collect(A::AbstractMatrix, B::SubArray) = dot(A, B)
_dot_collect(A::Diagonal, B::SubArray) = dot(A, collect(B))
end
22 changes: 22 additions & 0 deletions test/rulesets/LinearAlgebra/dense.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,4 +159,26 @@
test_rrule(lyap, A, C)
end
end
VERSION ≥ v"1.9.0" && @testset "kron" begin
@testset "AbstractVecOrMat{$T1}, AbstractVecOrMat{$T2}" for T1 in (Float64, ComplexF64), T2 in (Float64, ComplexF64)
@testset "frule" begin
test_frule(kron, randn(T1, 2), randn(T2, 3))
test_frule(kron, randn(T1, 2, 3), randn(T2, 5))
test_frule(kron, randn(T1, 2), randn(T2, 3, 5))
test_frule(kron, randn(T1, 2, 3), randn(T2, 5, 7))
end
@testset "rrule" begin
test_rrule(kron, randn(T1, 2), randn(T2, 3))

test_rrule(kron, Diagonal(randn(T1, 2)), randn(T2, 3))
test_rrule(kron, randn(T1, 2, 3), randn(T2, 5))

test_rrule(kron, randn(T1, 2), randn(T2, 3, 5))
test_rrule(kron, randn(T1, 2), Diagonal(randn(T2, 3)))

test_rrule(kron, randn(T1, 2, 3), randn(T2, 5, 7))
test_rrule(kron, Diagonal(randn(T1, 2)), Diagonal(randn(T2, 3)))
end
end
end
end