Using @F.functional decorator for forward method

Hi, while working on defining models using module_v3, there are some forward methods decorated with @F.functional. (example)

Upon some comments related to this, it seems that it allows possible fusion of operations that are executed in the method that is being decorated by making those ops atomic.

My question is

  1. Am I correct to understand its usage?
  2. If it’s right, does it only have an effect when using eager execution? In other words, does graph compiler automatically handle those fusions during compilation when using Graph API?

Hi! Great questions, we have some incoming documentation to help clarify as well :slight_smile: F.functional is probably doing too many things at once right now.

  1. Its core function is internal to the framework. It transforms functions which operate in the graph op space to functions which operate in the Tensor space. I hazard to say zero external framework users need this.
  2. It marks a function as “atomic” for eager execution, that is when you’re using functions eagerly that whole function will accumulate into a single graph, allowing graph optimizations like automatic fusion. This is occasionally useful to external users, but almost always is fine to not have. We should probably split it into a different function :slight_smile: for instance it’s really helpful on Linear.forwardbecause this function is particularly performance sensitive to fusion (hence why PyTorch has a custom kernel just for it, which we don’t need :smiley:)
  3. There’s some boring complexity we need to resolve with dimension algebra, you can’t do it outside a graph context today, so if you need to (say Dim(“x”) + 1) then it will need to be in a F.functional function for eager execution. We’ll fix this at some point, but it’s a pretty rare circumstance today.

does graph compiler automatically handle those fusions during compilation when using Graph API?

Yes, this will have no relevance when using Module.compile, or if you use any Tensoror functional methods with the Graph API!

It doesn’t exactly have no effect – it still will return Tensors instead of TensorValues, and the idea is that Tensors will be a bit nicer to work with even in graph authoring mode, we’ll tailor the API to be more familiar to model authors.

3 Likes

I agree that splitting it into a different function would be more intuitive for users. Thank you for the detailed explanation! I