Skip to content

Commit

Permalink
annotate type of operation field in Node type
Browse files Browse the repository at this point in the history
  • Loading branch information
ablaom committed Apr 8, 2024
1 parent 7ae5821 commit 190de70
Showing 1 changed file with 8 additions and 6 deletions.
14 changes: 8 additions & 6 deletions src/composition/learning_networks/nodes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,9 @@ See also [`node`](@ref), [`Source`](@ref), [`origins`](@ref),
[`sources`](@ref), [`fit!`](@ref).
"""
struct Node{T<:Union{Machine, Nothing}} <: AbstractNode
struct Node{T<:Union{Machine, Nothing},Oper} <: AbstractNode

operation # eg, `predict` or a static operation, such as `exp`
operation::Oper # eg, `predict` or a static operation, such as `exp`
machine::T # is `nothing` for static operations

# nodes called to get args for `operation(model, ...) ` or
Expand All @@ -43,9 +43,11 @@ struct Node{T<:Union{Machine, Nothing}} <: AbstractNode
# order consistent with extended graph, excluding self
nodes::Vector{AbstractNode}

function Node(operation,
machine::T,
args::AbstractNode...) where T<:Union{Machine, Nothing}
function Node(
operation::Oper,
machine::T,
args::AbstractNode...,
) where {T<:Union{Machine, Nothing}, Oper}

# check the number of arguments:
# if machine === nothing && isempty(args)
Expand All @@ -70,7 +72,7 @@ struct Node{T<:Union{Machine, Nothing}} <: AbstractNode
vcat(nodes_, (nodes(n) for n in machine.args)...) |> unique
end

return new{T}(operation, machine, args, origins_, nodes_)
return new{T,Oper}(operation, machine, args, origins_, nodes_)
end
end

Expand Down

0 comments on commit 190de70

Please sign in to comment.