Skip to content

Functional API: Execution environment agnostic function #205

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rnett opened this issue Feb 3, 2021 · 2 comments
Open

Functional API: Execution environment agnostic function #205

rnett opened this issue Feb 3, 2021 · 2 comments

Comments

@rnett
Copy link
Contributor

rnett commented Feb 3, 2021

cc @karllessard

This is somewhat of a sub-task of #181. The biggest pain I've ran into when using ConcreteFunction is that it only has tensor call methods, when it's mostly going to be used with Operand. This is a fairly simple issue on the surface. But, there's no way to execute ConcreteFunction in graph mode, i.e. if they are nested. The function used to generate the graph-mode outputs (which are wrapped in Signature) isn't saved. Now, it's easy enough to do this Java-side, in a sub class so that ConcreteFunction still supports loading. However, there's other issues such as supporting inputs with different shapes and dtypes that made me realize that what I'm trying to do here is closer to Python's Function and we may want to handle it with a new abstraction. There's also TF_Function and TF_GraphCopyFunction and TFE_ContextAddFunction which seems like it would allow attaching a ConcreteFunction to a graph without having to re-execute the builder in a new graph.

So I'd propose two things:

  • Implement Graph-mode and Eager-mode use of ConcreteFunctions using the native TF_Function APIs (the fact that the eager one doesn't mention gradients makes me a little worried, but I would think we can handle that manually later if necessary).
  • Add a Function class that acts like tf.function, in that it creates ConcreteFunctions as necessary for the argument shapes and dtypes. Additionally, since this will save the graph-creator lambda, we can have a debug flag that re-runs the lambda.

We also need to do something with variable handling, although that will probably need to wait on #179. Python seems to use an implicit variable-creation context to create them at the call-site and only allows it on the first call. I'd be fine with throwing errors and forcing the user to extract them, I think. I need to look into the details a bit more before I propose anything for this though. Variable scopes might be worth doing anyways for freezing, although hopefully explicit as part of Ops/Scope.

We'll need to pay attention to Graph states, like random seeds, too.

@rnett
Copy link
Contributor Author

rnett commented Feb 4, 2021

More thoughts about Function:

  • Inputs: Python auto-wraps everything in tensors, we can't do this. I'd like to allow non-tensor arguments, re-creating the graph if they change (treating them like shapes/dtypes). This also allows for some control flow structures inside the method, since we don't have access to autograph. I'd also like to include collections, and re-build if their structure changes (i.e. map keys, size of a list).
  • Outputs: I can't find any explicit listing of what Python allows, but it sounds like only things like Tensor or TensorList work. I'd limit ours to Operand and basic collections. They will have to be converted from the graph Operands to eager Operands, so they will be re-made, but doing that for List, Set, and Map at least isn't hard, and is worth doing to allow returning multiple things.
  • Closure variables: Python says it captures variables in it's closure, but it kind of lies: it only does this when it first builds the graph, updates to the variables won't be used. We could do something similar by auto-wrapping eager tensors in constantOf. I'd say we shouldn't and force use of constantOf, but then the behavior of the function changes whether it's executed in eager or graph mode (eager would work fine, graph would throw a runtime exception) if it's not used, so we probably should. The eager and graph behaviors still differ, but they do so in the same way Python's do. I would like a better way to handle this though, something like inputs.constantOf(() -> x) (see later section) may work.
  • Variables: My initial thinking here is that they need to be declared in an eager environment, then the resource tensors (which would need to be supported) passed to the graph as inputs. This would need to be hidden behind tf.Variable, probably by adding a VariableFactory or similar to Scope and using it. It would have to handle initialization too. Still need to look into it more, and given that we don't even have variables yet, support can be added later. It would be good to have a way to forbid creating tf.variable variable's though.
  • Construction: I have two modes in mind One is a lambda based one that takes a lambda like (tf, inputs) -> and gets/defines inputs using something like inputs.input("x"), which gets it in eager mode and declares a placeholder in graph mode. This might be better done through a subclass of Ops or a PlaceholderFactory like was mentioned for Variables. The other is a reflection based one that discovers inputs and outputs from the function. The biggest advantage of this is that it supports arbitrary signatures. Named arguments will likely be a problem though.

@rnett
Copy link
Contributor Author

rnett commented Feb 5, 2021

Note to myself: see how cross-graph references work wrt including graphs as functions, see #207.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant