Mill.jl/src/aggregations/segmented_max.jl at bd4b64902e67d5a2f05d096ec81fb9c97457ce9d · CTUAvastLab/Mill.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
"""
    SegmentedMax{V <: AbstractVector{<:Number}} <: AbstractAggregation

[`AbstractAggregation`](@ref) implementing segmented max aggregation:

``
f(\\{x_1, \\ldots, x_k\\}) = \\max_{i = 1, \\ldots, k} x_i
``

Stores a vector of parameters `ψ` that are filled into the resulting matrix in case an empty bag is encountered.

See also: [`AbstractAggregation`](@ref), [`AggregationStack`](@ref),
    [`SegmentedMean`](@ref), [`SegmentedSum`](@ref), [`SegmentedPNorm`](@ref), [`SegmentedLSE`](@ref).
"""
struct SegmentedMax{V <: AbstractVector{<:Number}} <: AbstractAggregation
    ψ::V
end

Flux.@layer :ignore SegmentedMax

SegmentedMax(T::Type, d::Integer) = SegmentedMax(zeros(T, d))
SegmentedMax(d::Integer) = SegmentedMax(Float32, d)

Flux.@forward SegmentedMax.ψ Base.getindex, Base.length, Base.size, Base.firstindex, Base.lastindex,
        Base.first, Base.last, Base.iterate, Base.eltype

Base.vcat(as::SegmentedMax...) = reduce(vcat, as |> collect)
function Base.reduce(::typeof(vcat), as::Vector{<:SegmentedMax})
    SegmentedMax(reduce(vcat, [a.ψ for a in as]))
end

function (a::SegmentedMax)(x::Maybe{AbstractMatrix{T}}, bags::AbstractBags,
                              w::Optional{AbstractVecOrMat{T}}=nothing) where T
    _check_agg(a, x)
    segmented_max_forw(x, a.ψ, bags)
end

segmented_max_forw(::Missing, ψ::AbstractVector, bags::AbstractBags) = repeat(ψ, 1, length(bags))
function segmented_max_forw(x::AbstractMatrix, ψ::AbstractVector, bags::AbstractBags)
    T = promote_type(eltype(x), eltype(ψ))
    y = Matrix{T}(fill(_typemin(T), size(x, 1), length(bags)))
    @inbounds for (bi, b) in enumerate(bags)
        if isempty(b)
            for i in eachindex(ψ)
                y[i, bi] = ψ[i]
            end
        else
            for j in b
                for i in axes(x, 1)
                    y[i, bi] = max(y[i, bi], x[i, j])
                end
            end
        end
    end
    y
end

function segmented_max_back(Δ, y, x, ψ, bags)
    dx = zero(x)
    dψ = zero(ψ)
    v = similar(x, size(x, 1))
    idxs = zeros(Int, size(x, 1))
    @inbounds for (bi, b) in enumerate(bags)
        if isempty(b)
            for i in eachindex(ψ)
                dψ[i] += Δ[i, bi]
            end
        else
            fi = first(b)
            v .= @view x[:, fi]
            idxs .= fi
            for j in b
                for i in axes(x, 1)
                    if v[i] < x[i, j]
                        idxs[i] = j
                        v[i] = x[i, j]
                    end
                end
            end
            for i in axes(x, 1)
                dx[i, idxs[i]] += Δ[i, bi]
            end
        end
    end
    dx, dψ, NoTangent()
end

function segmented_max_back(Δ, y, x::Missing, ψ, bags)
    dψ = zero(ψ)
    @inbounds for bi in eachindex(bags)
        for i in eachindex(ψ)
            dψ[i] += Δ[i, bi]
        end
    end
    ZeroTangent(), dψ, NoTangent()
end

function ChainRulesCore.rrule(::typeof(segmented_max_forw), args...)
    y = segmented_max_forw(args...)
    grad = Δ -> (NoTangent(), segmented_max_back(unthunk(Δ), y, args...)...)
    y, grad
end