Flux.jl/src/jit/inplace.jl
Mike J Innes 0e0057b0c4 basics
2018-02-09 13:51:07 +00:00

12 lines
255 B
Julia

mutable struct Cached{F,A}
f::F
buffer::A
end
function (c::Cached)(args...)
sh = shape(c.f, shape(args)...)
bytes(sh) > length(c.buffer) && (c.buffer = similar(c.buffer, bytes(sh)))
y = restructure(sh, c.buffer)
inplace!(c.f, y, args...)
end