TastyPancakes/Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

more general adaptors

Browse Source
This commit is contained in:
Mike J Innes 2017-10-04 18:55:56 +01:00
parent 2b95aff158
commit 1abc4febe6
2 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/onehot.jl
View File

@ -22,9 +22,12 @@ Base.:*(A::AbstractMatrix, B::OneHotMatrix) = A[:, map(x->x.ix, B.data)]
Base.hcat(x::OneHotVector, xs::OneHotVector...) = OneHotMatrix([x, xs...])
import NNlib.adapt
adapt(T, xs::OneHotMatrix) = OneHotMatrix(xs.height, adapt(T, xs.data))
@require CuArrays begin
import CuArrays: CuArray, cudaconvert
CuArrays.cu(xs::OneHotMatrix) = OneHotMatrix(xs.height, CuArrays.cu(xs.data))
Base.Broadcast._containertype(::Type{<:OneHotMatrix{<:CuArray}}) = CuArray
cudaconvert(x::OneHotMatrix{<:CuArray}) = OneHotMatrix(x.height, cudaconvert(x.data))
end

9
src/tracker/Tracker.jl
View File

@ -71,11 +71,10 @@ include("back.jl")
include("lib.jl")
include("numeric.jl")
using Requires
import NNlib.adapt
@require CuArrays begin
import CuArrays: cu
cu(xs::TrackedArray) = TrackedArray(xs.f, cu(xs.data), RefValue(cu(grad(xs))))
end
adapt(T, xs::TrackedArray) =
TrackedArray(xs.f, adapt(T, xs.data),
RefValue(adapt(T, grad(xs))))
end