[Data] Add API support for passing sample weights via ray.dataset.to_tf

[japneet-anyscale]

Description

Add support for ray.dataset.to_tf to pass sample weights

https://www.tensorflow.org/versions/r2.8/api_docs/python/tf/keras/Model#fit

API does not currently support the feature:
https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.to_tf.html#ray-data-dataset-to-tf

Use case

No response

[scottjlee]

@japneet-anyscale could you give more details / sample usage on how the new API would be used?

if i understand correctly, Model.fit(...) is for a Keras model, while the Dataset.to_tf(...) will produce a TF Dataset. So the two methods will be for different types (dataset vs. for fitting a model)

[scottjlee]

After internal discussion, got additional context. Current workaround is to create a new generator which yields the Dataset, and insert the sample_weights:

def generator():
                for batch in self.dataset.iter_batches(
                    batch_size=batch_size,
                ):
                    assert isinstance(batch, dict)
                    features = convert_batch_to_tensors(
                        batch, columns=feature_columns, type_spec=feature_type_spec
                    )
                    labels = convert_batch_to_tensors(
                        batch, columns=label_columns, type_spec=label_type_spec
                    )
                    sample_weights = convert_batch_to_tensors(
                        batch, columns=sample_weight_column, type_spec=sample_weight_spec
                    )
                    yield features, labels, sample_weights

[scottjlee]

We can add a new parameter additional_type_spec to to_tf():

def to_tf(
    ...,
    additional_type_spec: Dict[str, "tf.TypeSpec"] = None
)

then update generator() to output sample_weights, similar to the example in the previous message above: https://github.com/ray-project/ray/blob/releases/2.11.0/python/ray/data/iterator.py#L820-L835

finally, the API usage would be

to_tf(..., additional_type_spec={sample_weight_column: sample_weight_spec})

[galyna-anyscale]

@scottjlee Would you be able to pick this item in this sprint?

[scottjlee]

yeah either me or @omatthew98 can pick it up for the upcoming sprint (assuming in the next 6 weeks)

[jeffreyjeffreywang]

Hey @scottjlee, this issue looks interesting to me and I'm more than happy to help resolve this problem in the next couple days if you don't have spare cycles to tackle it. I am also thinking about adding tests to ray/python/ray/data/tests/test_tf.py as we introduce a new parameter to to_tf. Let me know if you prefer taking it over, or otherwise I'll publish a PR in the next couple days! Thanks!

[scottjlee]

@jeffreyjeffreywang that sounds great, go for it! thanks!