ray.data.range_tensor#

ray.data.range_tensor(n: int, *, shape: Tuple = (1,), parallelism: int = - 1) → ray.data.dataset.Dataset[ray.data._internal.arrow_block.ArrowRow][source]#

Create a Tensor dataset from a range of integers [0..n).

Examples

>>> import ray
>>> ds = ray.data.range_tensor(1000, shape=(2, 2)) 
>>> ds 
Dataset(
    num_blocks=200,
    num_rows=1000,
    schema={__value__: <ArrowTensorType: shape=(2, 2), dtype=int64>},
)
>>> ds.map_batches(lambda arr: arr * 2).take(2) 
[array([[0, 0],
        [0, 0]]),
array([[2, 2],
        [2, 2]])]

This is similar to range_table(), but uses the ArrowTensorArray extension type. The dataset elements take the form {“__value__”: array(N, shape=shape)}.

Parameters

  • n – The upper bound of the range of integer records.

  • shape – The shape of each record.

  • parallelism – The amount of parallelism to use for the dataset. Parallelism may be limited by the number of items.

Returns

Dataset holding the integers as Arrow tensor records.

PublicAPI: This API is stable across Ray releases.