Native functions v2

[rnett]

Replaces #211. Status is mostly the same, I still need to handle saving and gradients. It's also based on #232 instead of master, so the diff will get smaller once it's merged.

Currently the only ops forbidden in the body of functions is Placeholder. This could lead to somewhat weird situations when creating a function from a graph where you have something like:

val a = 1
val b = variable(2)

val c = a + c
val d = 10
val e = d + c

signature(inputs = (d), outputs = (e))

where even though e depends on c and d and c is not listed as an input, the function would work fine since c can be calculated without any inputs.

[rnett]

@karllessard my main issue now is that there is that the native functions don't work with our Init setup, which makes variables very hard to handle. You can't register init ops within a function or call them. It's necessary to get the variable behavior to match the old version (i.e. calling an incrementAndGet style ConcreteFunction with a variable inside it will increment it) and to export functions with variables. Note that since we only support exporting ConcreteFunction we need to handle variables somehow, or support exporting more things.

Optionally, we can just not allow variables in ConcreteFunction, as the @tf.function style API would support it properly by lifting initialization. This is my preferred solution, except we'll have to do something to allow training of imported models. Python seems to handle it by raising the variable declaration and initialization outside of the function with init_scope, which is essentially what I'm proposing here. I'm not sure how they handle exporting variables, though.

[rnett]

Another issue I need some guidance on: when loading a bundle, functions may call other functions that are part of the bundle (at least when loading the Python bundle). This doesn't work when the functions are extracted with function("name") and called, since there is no notion of function dependencies or including other functions, as far as I know. Is there any existing way around this? We could handle it Java side, but it would be a bit fragile.

[rnett]

After digging a bit more: the way we import things from Python (it might be the Python exporting, not sure if it's us), the native functions are called w/ placeholder inputs, and the imported MetaGraphSignatures use the wrappers. So the underlying functions are fine, and I can extract and use them, but that will differ slightly from the current export/import semantics (anything special in the call ops won't be included, I don't think there's that much?). Creating new functions from the wrappers causes the dependency issues.

Update: I can get the underlying functions from the call ops, but some are double wrapped and I can't unwrap that. I can get them via the signature names instead, which works, but is rather fragile.

[rnett]

I'm starting to think we shouldn't use ConcreteFunction for loading and instead have Tensor call methods on SavedModelBundle that make use of the loaded session. It's easier for working with pretrained models since variables and init will be handled correctly, and it removes a fair bit of overhead, and you can still use the ConcreteFunctions if needed. We could even use something like a FunctionView for calling a session w/ a preset signature.

[rnett]

Alright, dependencies are working. Of note, it correctly supports defining a ConcreteFunction that calls another function.

Also of note, we can not support ops with func type attributes using ConcreteFunction. I'll add the attribute methods in this PR, and save the actual ops for another one.

We could then use if for initialization, potentially, rather than the current method. It would work better w/ graph exporting and work in functions, too. Although I still want to look more at how Python handles it.

It would also be good to do statefulness disambiguation at the op level where possible, i.e. for If and StatelessIf. We can automatically call the stateless version if the passed function is staless. Detecting is somewhat more complicated.

[rnett]

Attaching a gradient function with a function does not appear to work, plus it causes issues with dependencies since it requires using the op name call style. PartitionedCall does not have c++ gradients defined yet. Since the original ConcreteFunction didn't support gradients either, I think we're ok adding it as-is and adding gradients to tensorflow core at some point. @karllessard thoughts?

[rnett]

The only issue remaining with this PR is initialization (variable exporting/importing). If we agree to delay that and rework initialization using if or similar, I think this is OK to merge. As things are it does not seem possible to handle "function side".

[rnett]

Gradients are working with tensorflow/tensorflow#47774

[rnett]

I also pushed op generation and the generated ops for ops that use functions. Eventually these should take a DeFun equivalent instead of ConcreteFunction (#238 (comment)), but they will work fine for now. Versions that take the function builder lambdas would be nice to have as well, and should be possible, but require more complicated codegen (and the new function API).

[karllessard]

Thanks @rnett ! Again, is it possible to break down this PR into smaller pieces? For instance, we can probably make the changes to ConcreteFunction before supporting the ops having a func argument, etc. That would simplify a lot the review process.

[rnett]

Yeah, I can pull out those changes. It'll get smaller once #232 gets merged, too.

[rnett]

Should be down to a reasonable size now. I just need to add the "function views" to SavedModelBundle if you're ok with handling things that way.

[karllessard]

Hey @rnett , just dropping a few comments so far but I've just started to review it, I'll add more comments soon

[karllessard]

Looks good @rnett ! Any chance you can revert the reformatting from your PR though? Only if it's easy for you to do.

[rnett]

Any chance of merging #308 first/soon? That's what I would use to do this.

[karllessard]

Any chance of merging #308 first/soon? That's what I would use to do this.

I guess we can merge it with these changes until we are set for a solution for reformatting. Can you rebase your PR though, there is a conflict there.

[rnett]

Also, could you merge #248 first (I'll rebase it in a moment)? Most of the changes are in here anyways, but it's the more up to date version.

[rnett]

Alright @karllessard, should be good to go.

[karllessard]

Hey @rnett , sorry for reviving an old debate but I'm wondering if we could play with the naming again... I know we are still in alpha but could be great to preserve some backward compatibility between our releases. For that, I would suggest the following name changes:

CallableFunction -> ConcreteFunction
ConcreteFunction -> GraphFunction

It make sense to me but does it to you? Basically a "concrete function" just mean a function that is "typed" and "realized" in Python, which make more or less sense in Java. Hence I'd be comfortable keeping it with my proposal. wdyt?

Doing this users won't need to do much when migrating from 0.3.1. Again, not a blocker, just a proposal.

[rnett]

That makes sense to me. Since I'm going to be adding more functional stuff, what do you think about moving all the function types to a function package? It's already starting to get cluttered.

[karllessard]

Do you think you can create a separate package without opening too many package-private methods?

[rnett]

Ugh, apparently not, there's too much use in Graph. Also, about the names, now that I think about it I don't like GraphFunction, since it works just fine with eager sessions too. I could use DefinedFunction, it will cause some of the higher level functions to not match Python, but I think that's fine, we're doing it here anyways.

[karllessard]

Also, about the names, now that I think about it I don't like GraphFunction, since it works just fine with eager sessions too.

But even in eager mode, the function created ends up being a graph that is executed "eagerly", no? So because of that, I think GraphFunction still seems adequate. And there is kind of a parallel with the SessionFunction that is issued from the saved models, where the function is always executed using a session while the other is added as a graph node.

I don't dislike DefinedFunction neither, just still want to keep the discussion opened for a bit.

[karllessard]

Awesome job @rnett , thanks!