[RFC] [Datasets] Extend to_tf to provide more flexibility

[bveeramani]

Description

This is a follow-up to #29028

Proposal

Add a parameter yield_fn to to_tf that lets user specify how elements are yielded.

Design considerations

What should yield_fn accept as input?

Option 1: yield_fn takes in the data normally yielded by to_tf

Inputs = tf.Tensor | list[tf.Tensor] | dict[str, tf.Tensor]
Targets = tf.Tensor | dict[str, tf.Tensor]

def to_tf(..., yield_fn: Callable[[Inputs, Targets], ...]):
    ...

Option 2: yield_fn takes in all of the data, and feature_colums and label_columns are ignored or passed to yield_fn

# Ignored case
def to_tf(..., yield_fn: [dict[str, tf.Tensor]], ...]):
    ...

# Passed to `yield_fn` case
def to_tf(..., yield_fn: [str | list[str], str | list[str], dict[str, tf.Tensor]], ...]):
    ...

What should yield_fn return as output?

Option 1: yield_fn returns a tuple of inputs and targets

Inputs = tf.Tensor | list[tf.Tensor] | dict[str, tf.Tensor]
Targets = tf.Tensor | dict[str, tf.Tensor]

def to_tf(..., yield_fn: Callable[[...], Tuple[Inputs, Targets]):

Option 2: yield_fn can return anything, but the user needs to specify the output signature

def to_tf(..., yield_fn: Callable[[...], Any], output_signature: Any):

Use case

Two potential use cases are:

1. Your model expects a list of tensors instead of a dict.

def yield_fn(inputs, targets):
    return list(inputs.values()), targets

2. Performing data manipulations that are difficult to express as UDFs. cc @jiaodong I think you have examples of this?

[amogkam]

Dumb question, but for use case 1, can't the user just convert a dict to a list before passing to the model?

For use case 2, isn't yield_fn already a UDF? What kind of UDFs would be easier to express via yield_fn?

[bveeramani]

Dumb question, but for use case 1, can't the user just convert a dict to a list before passing to the model?

I don’t think so. If you’re using Keras, you pass the TF Dataset created by to_tf directly to Model.fit. And, AFAIK, there’s no easy way to change the TF Datasets inputs to a list.

For use case 2, isn't yield_fn already a UDF?

Basically, except the inputs are dict[str, tf.Tensor. IIRC there are performance benefits to performing data manipulation in the from_generator level instead of a UDF.

[stale]

Hi, I'm a bot from the Ray team :)

To help human contributors to focus on more relevant issues, I will automatically add the stale label to issues that have had no activity for more than 4 months.

If there is no further activity in the 14 days, the issue will be closed!

• If you'd like to keep the issue open, just leave any comment, and the stale label will be removed!
• If you'd like to get more attention to the issue, please tag one of Ray's contributors.

You can always ask for help on our discussion forum or Ray's public slack channel.

[stale]

Hi again! The issue will be closed because there has been no more activity in the 14 days since the last message.

Please feel free to reopen or open a new issue if you'd still like it to be addressed.

Again, you can always ask for help on our discussion forum or Ray's public slack channel.

Thanks again for opening the issue!