Merge #1231
1231: use `ntuple` in conv r=MikeInnes a=MikeInnes This is the right abstraction over `map`, and in particular is a bit easier to compile away in some cases. As this is a trivial change from Flux's perspective it's not easy to test here, but there are downstream tests in XLA.jl. Co-authored-by: Mike J Innes <mike.j.innes@gmail.com>
This commit is contained in:
commit
254e4a7058
|
@ -132,7 +132,7 @@ end
|
|||
function (c::Conv)(x::AbstractArray)
|
||||
# TODO: breaks gpu broadcast :(
|
||||
# ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
|
||||
σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
|
||||
σ, b = c.σ, reshape(c.bias, ntuple(_->1, length(c.stride))..., :, 1)
|
||||
cdims = DenseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
|
||||
σ.(conv(x, c.weight, cdims) .+ b)
|
||||
end
|
||||
|
|
Loading…
Reference in New Issue