Flux.jl/src/tracker/Tracker.jl

module Tracker

export TrackedArray, TrackedVector, TrackedMatrix, param, back!

tracker(x) = nothing

istracked(x) = tracker(x) ≠ nothing
isleaf(x) = !istracked(x) || isleaf(tracker(x))
data(x) = istracked(x) ? data(tracker(x)) : x
grad(x) = grad(tracker(x))

struct Call{F,As<:Tuple}
  func::F
  args::As
end

Call(f, args...) = Call{typeof(f),typeof(args)}(f, args)

(c::Call)() = c.func(data.(c.args)...)

mutable struct Tracked{T}
  ref::UInt32
  f::Call
  data::T
  grad::T
  Tracked{T}(f::Call, data::T) where T = new(0, f, data)
  Tracked{T}(f::Call, data::T, grad::T) where T = new(0, f, data, grad)
end

Tracked(f::Call, x) = Tracked{typeof(x)}(f, x)

track(f::Call, x) = Tracked(f, x)
track(f::Call) = track(f, f())
track(f, xs...) = track(Call(f, xs...))

istracked(x::Tracked) = true
isleaf(x::Tracked) = x.f == Call(nothing)
data(x::Tracked) = x.data
grad(x::Tracked) = x.grad

include("back.jl")
include("scalar.jl")
include("array.jl")
include("numeric.jl")

param(x::Number) = TrackedNumber(float(x))
param(xs::AbstractArray) = TrackedArray(float.(xs))

using DataFlow
using DataFlow: inputnode, constant

vcall(f, args...) = vertex(DataFlow.Call(), constant(f), args...)
vcall(f::Broadcasted, args...) = vcall(broadcast, constant(f.f), args...)

_graph(x::Tracked, inputs...; cache = ObjectIdDict()) =
  vcall(x.f.func, map(x -> _graph(x, inputs...; cache = cache), x.f.args)...)

function _graph(x, inputs...; cache = ObjectIdDict())
  haskey(cache, x) && return cache[x]
  i = findfirst(inputs, x)
  cache[x] =
    i > 0 ? inputnode(i) :
    istracked(x) ? _graph(tracker(x), inputs...; cache = cache) :
    constant(x)
end

function graph(f, args...)
  inputs = param.(args)
  _graph(f(inputs...), inputs...)
end

import Adapt.adapt

adapt(T, xs::TrackedArray) = param(adapt(T, data(xs)))

end