[tf.data] Use a more efficient dispatch mechanism for functions in datasets.
This change adds an overload of the `FunctionLibraryRuntime::Run()` method that allows users to pass argument and return value containers in a `CallFrameInterface` object, rather than using the current (and expensive for large arities) default `FunctionCallFrame` implementation. It also specializes `CapturedFunction` to use this interface. Note that the new overload currently only supports local function execution, and more restructuring will be required to take advantage of it in the remote function execution case. This change should especially benefit datasets where each element has a large number of components (typically when training data have many features). PiperOrigin-RevId: 178684431
Loading
Please sign in to comment