ray.data.range_tensor#

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

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

Examples

>>> import ray
>>> ds = ray.data.range_tensor(1000, shape=(2, 2))
>>> ds  
Dataset(
   num_blocks=...,
   num_rows=1000,
   schema={data: numpy.ndarray(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 {“data”: 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 producing the integers as Arrow tensor records.

PublicAPI: This API is stable across Ray releases.