Some comments from a discussion

[oxinabox]

Was talking to @MikeInnes about this a while ago

Transcribing this here:

Lyndon White [10:41 AM]
Can you speak to this:
JuliaLogging/TensorBoardLogger.jl#32 (comment)
My feeling is it is AI-Hard, to decide from IR what is a object that should be in the graph.
But there might be some clever hacks like defining things that have sensitivities defined as fundermental objects

I also don't think it is useful, but maybe I am wrong.

Mike Innes [10:29 AM]
Yeah, I wonder if you could do a cassette-style pass that partially evaluates some things and holds back others
Of course you’d have to recover control flow (assuming the visualiser can support that?)
But yes, in general you’re trying to visualise an arbitrary Julia program, which doesn’t really work

So I suspect what you really want here is to visualise the tree structure of the model itself, not the graph structure of the execution of the model, and for something like resnet where this is actually useful they should be roughly the same

Lyndon White [10:41 AM]
Maybe we can consume what @treelike creates?

Mike Innes [10:43 AM]
Yes, exactly
You just need to recurse over children

[shashikdm]

@treelike seems to overload children for the input struct and map its constructor to fieldnames. I don't understand how it can be used to build a graph. please explain how. Also, it may not be able to consider functions such as x->Tracker.data(x)

[MikeInnes]

The recursive set of children is a graph (well, a tree). One simple thing to do here is to make the graph relationship between the children overloadable, so that Chain and SkipConnection have the right graphs and ResNet comes out looking reasonable.

It's true that this can't handle anonymous functions, but neither can a method that looks at IR in general. Julia is a dynamic language and any dynamic dispatch is going to thwart a static analysis.

[oxinabox]

Does @treelike overload AbstractTrees.children?
That would be nice for purposes of avoiding a direct Flux dependency here,
and hopefully ending up with something more general.

[MikeInnes]

No, it's Flux.children. In the past we have discussed splitting it out. I'm a bit wary because it's not clear that it isn't essentially a Flux-specific abstraction at the moment. But I guess we could call it FluxTrees.jl or something.

[oxinabox]

Creating a package called FluxTrees.jl does not seems sensible.

[MikeInnes]

Is that an objection to the package or the specific name? We can bikeshed the name easily enough if avoiding the flux dependency seems useful enough.

[oxinabox]

breaking a single function into a package, seems excessive.
Why not just overload AbstractTrees?

[MikeInnes]

We could try that, but I don't think they have the same semantics. Flux doesn't really want AbstractTrees provides and I don't think AbstractTrees want what Flux provides. I'm not sure they're compatible if you e.g. put a tree inside a Flux model. This is why I say that Flux's current abstraction seems Flux-specific (though of course designing something that isn't Flux specific to do the same job would be great).

[oxinabox]

It would be good the have an abstraction that works for both Flux and Turing.jl

[shashikdm]

Please checkout dev branch. IR approach seems promising to me. I haven't written cases for all kinds of Expr in IR but I believe that will not be very difficult. Also I am yet to add more function names to norecurselist. That is also doable. Download (https://github.com/shashikdm/TraceGraph.jl/blob/assets/events.out.tfevents.1.560585414358108e9.shashi-Aspire-A515-51G) and run tensorboardlogger with it to interact with the graph for

function foo(a)
  b = a+10
  c = b+10
  return (a,b,c)
end

tg = tracegraph(foo, 10)

What is your opinion on this approach?
Do you think that after adding more cases to cover all kinds of Expr and populating norecurselist this will work?

[shashikdm]

this is the gplot of tracegraph:

this is tensorboardgraph:

[shashikdm]

It's true that this can't handle anonymous functions, but neither can a method that looks at IR in general. Julia is a dynamic language and any dynamic dispatch is going to thwart a static analysis.

I'm trying to do dynamic analysis by taking function and args using Cassette. It should be able to handle anonymous functions too.

[PhilipVinc]

Wow, very nice progress!