ray.data.range_tensor
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.