pyarrow.fixed_shape_tensor

pyarrow.fixed_shape_tensor(DataType value_type, shape, dim_names=None, permutation=None)

Create instance of fixed shape tensor extension type with shape and optional names of tensor dimensions and indices of the desired logical ordering of dimensions.

Parameters:

value_typeDataType

Data type of individual tensor elements.

shapetuple or list of integers

The physical shape of the contained tensors.

dim_namestuple or list of strings, default None

Explicit names to tensor dimensions.

permutationtuple or list integers, default None

Indices of the desired ordering of the original dimensions. The indices contain a permutation of the values [0, 1, .., N-1] where N is the number of dimensions. The permutation indicates which dimension of the logical layout corresponds to which dimension of the physical tensor. For more information on this parameter see Fixed shape tensor.

Returns:

typeFixedShapeTensorType

Examples

Create an instance of fixed shape tensor extension type:

>>> import pyarrow as pa
>>> tensor_type = pa.fixed_shape_tensor(pa.int32(), [2, 2])
>>> tensor_type
FixedShapeTensorType(extension<arrow.fixed_shape_tensor[value_type=int32, shape=[2,2]]>)

Inspect the data type:

>>> tensor_type.value_type
DataType(int32)
>>> tensor_type.shape
[2, 2]

Create a table with fixed shape tensor extension array:

>>> arr = [[1, 2, 3, 4], [10, 20, 30, 40], [100, 200, 300, 400]]
>>> storage = pa.array(arr, pa.list_(pa.int32(), 4))
>>> tensor = pa.ExtensionArray.from_storage(tensor_type, storage)
>>> pa.table([tensor], names=["tensor_array"])
pyarrow.Table
tensor_array: extension<arrow.fixed_shape_tensor[value_type=int32, shape=[2,2]]>
----
tensor_array: [[[1,2,3,4],[10,20,30,40],[100,200,300,400]]]

Create an instance of fixed shape tensor extension type with names of tensor dimensions:

>>> tensor_type = pa.fixed_shape_tensor(pa.int8(), (2, 2, 3),
...                                     dim_names=['C', 'H', 'W'])
>>> tensor_type.dim_names
['C', 'H', 'W']

Create an instance of fixed shape tensor extension type with permutation:

>>> tensor_type = pa.fixed_shape_tensor(pa.int8(), (2, 2, 3),
...                                     permutation=[0, 2, 1])
>>> tensor_type.permutation
[0, 2, 1]