Feature request - return better types from `batch`

[oveddan]

Currently the method signature for batch always returns a Dataset<TensorContainer> which can be challenging to use again when you want to dig into the batched dataset as you have to cast the type.

For example:

const data = tf.data.array([1, 2, 3, 4, 5]);

// sequences is a Dataset<TensorContainer>
const sequences = data.batch(2, false);

const ioDataset = sequences.map(x => {
   const asTensor = x as Tensor1D;
   
    const inputText = asTensor.slice([0], [asTensor.shape[0] - 1]);
    const outputText = asTensor.slice([1], [asTensor.shape[0] - 1]);

    return [inputText, outputText];
});

It would be nice if the batch function above returned a Dataset<Tensor1D>

[davidsoergel]

Good point, thanks! I think this is possible but there may be some tricky edge cases-- we'll take a look.

[davidsoergel]

And similarly for take() and all the rest :)

[gaikwadrahul8]

Hi, @oveddan

Apologize for the delayed response and I see this issue is too old so May I know have we taken care of this feature request or not ? If possible could you please try to replicate this issue with latest version of@tensorflow/[email protected] ? Thank you!

[google-ml-butler[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs. Thank you.

[google-ml-butler[bot]]

Closing as stale. Please @mention us if this needs more attention.