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