-
Notifications
You must be signed in to change notification settings - Fork 137
Adds a mechanism to manually promote constants to inputs in Lazy Tensor. #299
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks fine. Could you include a motivating example in the description?
+1 for documenting a motivating example. Also, could we start with this being |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If this protocol is internal, it need not be underscored. The underscoring rule only applies to implementation-detail types that had to be made public.
|
||
extension TensorHandle: _LazyTensorCompatible { | ||
var _lazyTensor: LazyTensor? { handle._lazyTensor } | ||
public var _concreteLazyTensor: TensorHandle { TensorHandle(handle: handle._concreteLazyTensor) } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The protocol is internal. Why does this requirement need to be public? Same for other occurrences of public
.
- Only add _concreteInputLazyTensor to Tensor types. - Move _LazyTensorCompatible protocol to a test helper.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Constants in the example don't match up.
Thanks for flagging it. Fixed now. |
extension TensorHandle { | ||
/// Returns `Self` that wraps `_concreteInputLazyTensor` of the underlying | ||
/// `_AnyTensorHandle` | ||
public var _concreteInputLazyTensor: TensorHandle { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm is there a reason this needs to be public?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I expect that it will only be used within the library. However, I wanted it to be available for experimentation in the user models if needed.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
How and why should a user be expected to use this?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just wanted to be careful exposing public APIs. I think we need to make sure public APIs have at least some justification of use cases before they are exposed.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The use case is the one in description where we want to force a constant to be an input to an extracted trace. Ultimately, we may not need this API. However, in the meantime, I want to have this knob while we develop and improve the lazy tensor implementation.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh I see. Sorry for missing parts of the PR description! And thanks for explaining :)
Adds a
_concreteInputLazyTensor
property to the tensor-related types in the library. This provides temporary mechanism to promote constants to inputs manually in trace. Here is an example use case:This will result in the following trace for
w
:Note that the slot for
a
got promoted to an input. This is added to experiment with the lazy tensor implementation and is mostly intended to be used within theTensorFlow
module. If this turns out to be useful, we should provide a cleaner user-visible API.This PR also adds a helper file for tests related to lazy tensor.