Skip to content
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

Jump to bottom

Support conv_direct! on custom datatypes #592

Merged
merged 4 commits into from
Jun 19, 2024
Merged

Support conv_direct! on custom datatypes #592

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3336842
Add test for unusual input datatypes
adrhill Jun 17, 2024
0e318fd
Add fix to `conv_direct!`
adrhill Jun 17, 2024
68e15a9
Only set y to zero if beta is false or zero
adrhill Jun 17, 2024
8f87596
Test output eltype
adrhill Jun 17, 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
5 changes: 5 additions & 0 deletions src/impl/conv_direct.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,11 @@ function conv_direct!(
# Use `calc_padding_regions` to determine where we do or don't need to worry about padding
padded_regions, central_region = calc_padding_regions(cdims)

# Set outputs to zero to support custom datatypes (https://github.com/FluxML/NNlib.jl/issues/490)
if iszero(beta)
adrhill marked this conversation as resolved.
Show resolved Hide resolved
y = fill!(y, zero(yT))
Copy link

Choose a reason for hiding this comment

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

As I mentioned in the original issue, naive fill!/fill doesn't necessarily work with e.g. BigInt:

julia> a = Vector{BigInt}(undef, 2)
2-element Vector{BigInt}:
 #undef
 #undef

julia> fill!(a, zero(BigInt))
2-element Vector{BigInt}:
 0
 0

julia> a[1] ===a[2]
true

I don't know how important these kinds of usecases are for NNlib.jl though.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Is that problematic here?

julia> a = Vector{BigInt}(undef, 2)
2-element Vector{BigInt}:
 #undef
 #undef

julia> fill!(a, zero(BigInt))
2-element Vector{BigInt}:
 0
 0

julia> for (i, ai) in enumerate(a)
         a[i] = i + 0 * a[i]
       end

julia> a
2-element Vector{BigInt}:
 1
 2

end

# Start with the central region
w_region, h_region, d_region = central_region
@inbounds for batch in 1:size(x, 5),
Expand Down
39 changes: 39 additions & 0 deletions test/conv.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ using NNlib, Test
using NNlib: input_size, kernel_size, channels_in, channels_out, channel_multiplier,
stride, padding, dilation, flipkernel, output_size,
groupcount
using Random: AbstractRNG, SamplerType

@testset "ConvDims" begin
for T in (DenseConvDims, DepthwiseConvDims)
Expand Down Expand Up @@ -865,6 +866,44 @@ end
@test size(NNlib.∇conv_filter_direct!(w, x, y, cdims)) == w_size
end

# https://github.com/FluxML/NNlib.jl/issues/490
# https://github.com/FluxML/NNlib.jl/issues/405
@testset "conv_direct! - Unusual input types" begin
# Create test type that can't be indexed when undefined.
# This simulates the worst-case scenario for custom types.
struct MyFloat <: Real
set::Set{Float32}
end

# Test that direct indexing fails when undefined.
v = Array{MyFloat}(undef, 3)
@test_throws UndefRefError v[1]

# Define minimal set of functions required for conv_direct!
MyFloat(x::MyFloat) = x
MyFloat(x::Real) = MyFloat(Set(Float32(x)))

Base.:+(x::MyFloat, y::MyFloat) = MyFloat(only(x.set) + only(y.set))
Base.:*(x::MyFloat, y::MyFloat) = MyFloat(only(x.set) * only(y.set))
Base.promote_rule(::Type{MyFloat}, ::Type{Float32}) = MyFloat
Base.rand(::AbstractRNG, ::SamplerType{MyFloat}) = MyFloat(rand(Float32))
Base.zero(::MyFloat) = MyFloat(zero(Float32))
Base.zero(::Type{MyFloat}) = MyFloat(zero(Float32))

# Test conv_direct!
x_size = (6, 7, 8, 5, 3)
y_size = (5, 6, 7, 4, 3)
w_size = (2, 2, 2, 5, 4)
x = rand(MyFloat, x_size);
w = randn(Float32, w_size);
y = Array{MyFloat}(undef, y_size...);
cdims = DenseConvDims(x_size, w_size)
y_out = NNlib.conv_direct!(y, x, w, cdims)

@test eltype(y_out) == MyFloat
@test size(y_out) == y_size
end

@testset "AutoDiff: spatial_rank=$spatial_rank" for spatial_rank in (1, 2, 3)
x = rand(rng, repeat([5], spatial_rank)..., 3, 2)
w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
Expand Down
Loading