diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index af8305a025291..02e0a890e4b3d 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -19,6 +19,8 @@ #include #include "arrow/extension/fixed_shape_tensor.h" +#include "arrow/extension/tensor_internal.h" +#include "arrow/scalar.h" #include "arrow/array/array_nested.h" #include "arrow/array/array_primitive.h" @@ -86,7 +88,7 @@ bool FixedShapeTensorType::ExtensionEquals(const ExtensionType& other) const { if (extension_name() != other.extension_name()) { return false; } - const auto& other_ext = static_cast(other); + const auto& other_ext = internal::checked_cast(other); auto is_permutation_trivial = [](const std::vector& permutation) { for (size_t i = 1; i < permutation.size(); ++i) { @@ -143,7 +145,7 @@ std::string FixedShapeTensorType::Serialize() const { if (!dim_names_.empty()) { rj::Value dim_names(rj::kArrayType); - for (std::string v : dim_names_) { + for (const std::string& v : dim_names_) { dim_names.PushBack(rj::Value{}.SetString(v.c_str(), allocator), allocator); } document.AddMember(rj::Value("dim_names", allocator), dim_names, allocator); @@ -199,10 +201,52 @@ std::shared_ptr FixedShapeTensorType::MakeArray( std::shared_ptr data) const { DCHECK_EQ(data->type->id(), Type::EXTENSION); DCHECK_EQ("arrow.fixed_shape_tensor", - static_cast(*data->type).extension_name()); + internal::checked_cast(*data->type).extension_name()); return std::make_shared(data); } +Result> FixedShapeTensorType::MakeTensor( + const std::shared_ptr& scalar) { + const auto ext_scalar = internal::checked_pointer_cast(scalar); + const auto ext_type = + internal::checked_pointer_cast(scalar->type); + if (!is_fixed_width(*ext_type->value_type())) { + return Status::TypeError("Cannot convert non-fixed-width values to Tensor."); + } + const auto array = + internal::checked_pointer_cast(ext_scalar->value)->value; + if (array->null_count() > 0) { + return Status::Invalid("Cannot convert data with nulls to Tensor."); + } + const auto value_type = + internal::checked_pointer_cast(ext_type->value_type()); + const auto byte_width = value_type->byte_width(); + + std::vector permutation = ext_type->permutation(); + if (permutation.empty()) { + permutation.resize(ext_type->ndim()); + std::iota(permutation.begin(), permutation.end(), 0); + } + + std::vector shape = ext_type->shape(); + internal::Permute(permutation, &shape); + + std::vector dim_names = ext_type->dim_names(); + if (!dim_names.empty()) { + internal::Permute(permutation, &dim_names); + } + + std::vector strides; + RETURN_NOT_OK(ComputeStrides(*value_type.get(), shape, permutation, &strides)); + const auto start_position = array->offset() * byte_width; + const auto size = std::accumulate(shape.begin(), shape.end(), static_cast(1), + std::multiplies<>()); + const auto buffer = + SliceBuffer(array->data()->buffers[1], start_position, size * byte_width); + + return Tensor::Make(ext_type->value_type(), buffer, shape, strides, dim_names); +} + Result> FixedShapeTensorArray::FromTensor( const std::shared_ptr& tensor) { auto permutation = internal::ArgSort(tensor->strides(), std::greater<>()); @@ -293,53 +337,71 @@ const Result> FixedShapeTensorArray::ToTensor() const { // To convert an array of n dimensional tensors to a n+1 dimensional tensor we // interpret the array's length as the first dimension the new tensor. - auto ext_arr = std::static_pointer_cast(this->storage()); - auto ext_type = internal::checked_pointer_cast(this->type()); - ARROW_RETURN_IF(!is_fixed_width(*ext_arr->value_type()), - Status::Invalid(ext_arr->value_type()->ToString(), - " is not valid data type for a tensor")); - auto permutation = ext_type->permutation(); - - std::vector dim_names; - if (!ext_type->dim_names().empty()) { - for (auto i : permutation) { - dim_names.emplace_back(ext_type->dim_names()[i]); - } - dim_names.insert(dim_names.begin(), 1, ""); + const auto ext_type = + internal::checked_pointer_cast(this->type()); + const auto value_type = ext_type->value_type(); + ARROW_RETURN_IF( + !is_fixed_width(*value_type), + Status::TypeError(value_type->ToString(), " is not valid data type for a tensor")); + + // ext_type->permutation() gives us permutation for a single row with values in + // range [0, ndim). Here want to create a ndim + 1 dimensional tensor from the entire + // array and we assume the first dimension will always have the greatest stride, so it + // will get permutation index 0 and remaining values from ext_type->permutation() need + // to be shifted to fill the [1, ndim+1) range. Computed permutation will be used to + // generate the new tensor's shape, strides and dim_names. + std::vector permutation = ext_type->permutation(); + if (permutation.empty()) { + permutation.resize(ext_type->ndim() + 1); + std::iota(permutation.begin(), permutation.end(), 0); } else { - dim_names = {}; + for (auto i = 0; i < static_cast(ext_type->ndim()); i++) { + permutation[i] += 1; + } + permutation.insert(permutation.begin(), 1, 0); } - std::vector shape; - for (int64_t& i : permutation) { - shape.emplace_back(ext_type->shape()[i]); - ++i; + std::vector dim_names = ext_type->dim_names(); + if (!dim_names.empty()) { + dim_names.insert(dim_names.begin(), 1, ""); + internal::Permute(permutation, &dim_names); } + + std::vector shape = ext_type->shape(); + auto cell_size = std::accumulate(shape.begin(), shape.end(), static_cast(1), + std::multiplies<>()); shape.insert(shape.begin(), 1, this->length()); - permutation.insert(permutation.begin(), 1, 0); + internal::Permute(permutation, &shape); std::vector tensor_strides; - auto value_type = internal::checked_pointer_cast(ext_arr->value_type()); + const auto fw_value_type = internal::checked_pointer_cast(value_type); ARROW_RETURN_NOT_OK( - ComputeStrides(*value_type.get(), shape, permutation, &tensor_strides)); - ARROW_ASSIGN_OR_RAISE(auto buffers, ext_arr->Flatten()); + ComputeStrides(*fw_value_type.get(), shape, permutation, &tensor_strides)); + + const auto raw_buffer = this->storage()->data()->child_data[0]->buffers[1]; ARROW_ASSIGN_OR_RAISE( - auto tensor, Tensor::Make(ext_arr->value_type(), buffers->data()->buffers[1], shape, - tensor_strides, dim_names)); - return tensor; + const auto buffer, + SliceBufferSafe(raw_buffer, this->offset() * cell_size * value_type->byte_width())); + + return Tensor::Make(value_type, buffer, shape, tensor_strides, dim_names); } Result> FixedShapeTensorType::Make( const std::shared_ptr& value_type, const std::vector& shape, const std::vector& permutation, const std::vector& dim_names) { - if (!permutation.empty() && shape.size() != permutation.size()) { - return Status::Invalid("permutation size must match shape size. Expected: ", - shape.size(), " Got: ", permutation.size()); + const auto ndim = shape.size(); + if (!permutation.empty() && ndim != permutation.size()) { + return Status::Invalid("permutation size must match shape size. Expected: ", ndim, + " Got: ", permutation.size()); + } + if (!dim_names.empty() && ndim != dim_names.size()) { + return Status::Invalid("dim_names size must match shape size. Expected: ", ndim, + " Got: ", dim_names.size()); } - if (!dim_names.empty() && shape.size() != dim_names.size()) { - return Status::Invalid("dim_names size must match shape size. Expected: ", - shape.size(), " Got: ", dim_names.size()); + if (!permutation.empty()) { + RETURN_NOT_OK(internal::IsPermutationValid(permutation)); } + const auto size = std::accumulate(shape.begin(), shape.end(), static_cast(1), std::multiplies<>()); return std::make_shared(value_type, static_cast(size), diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.h b/cpp/src/arrow/extension/fixed_shape_tensor.h index fcfb1ebbab96a..591a7cee32a34 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.h +++ b/cpp/src/arrow/extension/fixed_shape_tensor.h @@ -64,7 +64,7 @@ class ARROW_EXPORT FixedShapeTensorType : public ExtensionType { std::string ToString() const override; /// Number of dimensions of tensor elements - size_t ndim() { return shape_.size(); } + size_t ndim() const { return shape_.size(); } /// Shape of tensor elements const std::vector shape() const { return shape_; } @@ -94,6 +94,15 @@ class ARROW_EXPORT FixedShapeTensorType : public ExtensionType { /// Create a FixedShapeTensorArray from ArrayData std::shared_ptr MakeArray(std::shared_ptr data) const override; + /// \brief Create a Tensor from an ExtensionScalar from a FixedShapeTensorArray + /// + /// This method will return a Tensor from ExtensionScalar with strides + /// derived from shape and permutation of FixedShapeTensorType. Shape and + /// dim_names will be permuted according to permutation stored in the + /// FixedShapeTensorType metadata. + static Result> MakeTensor( + const std::shared_ptr& scalar); + /// \brief Create a FixedShapeTensorType instance static Result> Make( const std::shared_ptr& value_type, const std::vector& shape, diff --git a/cpp/src/arrow/extension/fixed_shape_tensor_test.cc b/cpp/src/arrow/extension/fixed_shape_tensor_test.cc index 2b8e703d3c66e..3fd39a11ff50d 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor_test.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor_test.cc @@ -28,6 +28,7 @@ #include "arrow/tensor.h" #include "arrow/testing/gtest_util.h" #include "arrow/util/key_value_metadata.h" +#include "arrow/util/sort.h" namespace arrow { @@ -39,34 +40,34 @@ class TestExtensionType : public ::testing::Test { public: void SetUp() override { shape_ = {3, 3, 4}; - cell_shape_ = {3, 4}; + element_shape_ = {3, 4}; value_type_ = int64(); - cell_type_ = fixed_size_list(value_type_, 12); + element_type_ = fixed_size_list(value_type_, 12); dim_names_ = {"x", "y"}; ext_type_ = internal::checked_pointer_cast( - fixed_shape_tensor(value_type_, cell_shape_, {}, dim_names_)); + fixed_shape_tensor(value_type_, element_shape_, {}, dim_names_)); values_ = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35}; values_partial_ = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}; shape_partial_ = {2, 3, 4}; tensor_strides_ = {96, 32, 8}; - cell_strides_ = {32, 8}; + element_strides_ = {32, 8}; serialized_ = R"({"shape":[3,4],"dim_names":["x","y"]})"; } protected: std::vector shape_; std::vector shape_partial_; - std::vector cell_shape_; + std::vector element_shape_; std::shared_ptr value_type_; - std::shared_ptr cell_type_; + std::shared_ptr element_type_; std::vector dim_names_; std::shared_ptr ext_type_; std::vector values_; std::vector values_partial_; std::vector tensor_strides_; - std::vector cell_strides_; + std::vector element_strides_; std::string serialized_; }; @@ -96,8 +97,8 @@ TEST_F(TestExtensionType, CreateExtensionType) { // Test ExtensionType methods ASSERT_EQ(ext_type_->extension_name(), "arrow.fixed_shape_tensor"); ASSERT_TRUE(ext_type_->Equals(*exact_ext_type)); - ASSERT_FALSE(ext_type_->Equals(*cell_type_)); - ASSERT_TRUE(ext_type_->storage_type()->Equals(*cell_type_)); + ASSERT_FALSE(ext_type_->Equals(*element_type_)); + ASSERT_TRUE(ext_type_->storage_type()->Equals(*element_type_)); ASSERT_EQ(ext_type_->Serialize(), serialized_); ASSERT_OK_AND_ASSIGN(auto ds, ext_type_->Deserialize(ext_type_->storage_type(), serialized_)); @@ -106,18 +107,28 @@ TEST_F(TestExtensionType, CreateExtensionType) { // Test FixedShapeTensorType methods ASSERT_EQ(exact_ext_type->id(), Type::EXTENSION); - ASSERT_EQ(exact_ext_type->ndim(), cell_shape_.size()); - ASSERT_EQ(exact_ext_type->shape(), cell_shape_); + ASSERT_EQ(exact_ext_type->ndim(), element_shape_.size()); + ASSERT_EQ(exact_ext_type->shape(), element_shape_); ASSERT_EQ(exact_ext_type->value_type(), value_type_); - ASSERT_EQ(exact_ext_type->strides(), cell_strides_); + ASSERT_EQ(exact_ext_type->strides(), element_strides_); ASSERT_EQ(exact_ext_type->dim_names(), dim_names_); EXPECT_RAISES_WITH_MESSAGE_THAT( Invalid, testing::HasSubstr("Invalid: permutation size must match shape size."), - FixedShapeTensorType::Make(value_type_, cell_shape_, {0})); + FixedShapeTensorType::Make(value_type_, element_shape_, {0})); EXPECT_RAISES_WITH_MESSAGE_THAT( Invalid, testing::HasSubstr("Invalid: dim_names size must match shape size."), - FixedShapeTensorType::Make(value_type_, cell_shape_, {}, {"x"})); + FixedShapeTensorType::Make(value_type_, element_shape_, {}, {"x"})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + testing::HasSubstr("Invalid: Permutation indices for 2 dimensional tensors must be " + "unique and within [0, 1] range. Got: [3,0]"), + FixedShapeTensorType::Make(value_type_, {5, 6}, {3, 0})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + testing::HasSubstr("Invalid: Permutation indices for 3 dimensional tensors must be " + "unique and within [0, 2] range. Got: [0,1,1]"), + FixedShapeTensorType::Make(value_type_, {1, 2, 3}, {0, 1, 1})); } TEST_F(TestExtensionType, EqualsCases) { @@ -148,7 +159,7 @@ TEST_F(TestExtensionType, CreateFromArray) { std::vector> buffers = {nullptr, Buffer::Wrap(values_)}; auto arr_data = std::make_shared(value_type_, values_.size(), buffers, 0, 0); auto arr = std::make_shared(arr_data); - ASSERT_OK_AND_ASSIGN(auto fsla_arr, FixedSizeListArray::FromArrays(arr, cell_type_)); + ASSERT_OK_AND_ASSIGN(auto fsla_arr, FixedSizeListArray::FromArrays(arr, element_type_)); auto ext_arr = ExtensionType::WrapArray(ext_type_, fsla_arr); ASSERT_EQ(ext_arr->length(), shape_[0]); ASSERT_EQ(ext_arr->null_count(), 0); @@ -200,7 +211,7 @@ TEST_F(TestExtensionType, RoundtripBatch) { std::vector> buffers = {nullptr, Buffer::Wrap(values_)}; auto arr_data = std::make_shared(value_type_, values_.size(), buffers, 0, 0); auto arr = std::make_shared(arr_data); - ASSERT_OK_AND_ASSIGN(auto fsla_arr, FixedSizeListArray::FromArrays(arr, cell_type_)); + ASSERT_OK_AND_ASSIGN(auto fsla_arr, FixedSizeListArray::FromArrays(arr, element_type_)); auto ext_arr = ExtensionType::WrapArray(ext_type_, fsla_arr); // Pass extension array, expect getting back extension array @@ -215,7 +226,7 @@ TEST_F(TestExtensionType, RoundtripBatch) { auto ext_metadata = key_value_metadata({{"ARROW:extension:name", exact_ext_type->extension_name()}, {"ARROW:extension:metadata", serialized_}}); - ext_field = field(/*name=*/"f0", /*type=*/cell_type_, /*nullable=*/true, + ext_field = field(/*name=*/"f0", /*type=*/element_type_, /*nullable=*/true, /*metadata=*/ext_metadata); auto batch2 = RecordBatch::Make(schema({ext_field}), fsla_arr->length(), {fsla_arr}); RoundtripBatch(batch2, &read_batch2); @@ -270,7 +281,7 @@ TEST_F(TestExtensionType, CreateFromTensor) { auto ext_arr_5 = std::static_pointer_cast( ExtensionType::WrapArray(ext_type_5, fsla_arr)); EXPECT_RAISES_WITH_MESSAGE_THAT( - Invalid, testing::HasSubstr("binary is not valid data type for a tensor"), + TypeError, testing::HasSubstr("binary is not valid data type for a tensor"), ext_arr_5->ToTensor()); auto ext_type_6 = internal::checked_pointer_cast( @@ -278,6 +289,10 @@ TEST_F(TestExtensionType, CreateFromTensor) { auto arr_with_null = ArrayFromJSON(int64(), "[1, 0, null, null, 1, 2]"); ASSERT_OK_AND_ASSIGN(auto fsla_arr_6, FixedSizeListArray::FromArrays( arr_with_null, fixed_size_list(int64(), 2))); + + auto ext_type_7 = internal::checked_pointer_cast( + fixed_shape_tensor(int64(), {3, 4}, {})); + ASSERT_OK_AND_ASSIGN(auto ext_arr_7, FixedShapeTensorArray::FromTensor(tensor)); } void CheckFromTensorType(const std::shared_ptr& tensor, @@ -308,7 +323,7 @@ TEST_F(TestExtensionType, TestFromTensorType) { auto dim_names = std::vector>{ {"y", "z"}, {"z", "y"}, {"y", "z"}, {"z", "y"}, {"y", "z"}, {"y", "z"}, {"y", "z"}, {"y", "z"}}; - auto cell_shapes = std::vector>{{3, 4}, {4, 3}, {4, 3}, {3, 4}}; + auto element_shapes = std::vector>{{3, 4}, {4, 3}, {4, 3}, {3, 4}}; auto permutations = std::vector>{{0, 1}, {1, 0}, {0, 1}, {1, 0}}; for (size_t i = 0; i < shapes.size(); i++) { @@ -316,11 +331,82 @@ TEST_F(TestExtensionType, TestFromTensorType) { strides[i], tensor_dim_names[i])); ASSERT_OK_AND_ASSIGN(auto ext_arr, FixedShapeTensorArray::FromTensor(tensor)); auto ext_type = - fixed_shape_tensor(value_type_, cell_shapes[i], permutations[i], dim_names[i]); + fixed_shape_tensor(value_type_, element_shapes[i], permutations[i], dim_names[i]); CheckFromTensorType(tensor, ext_type); } } +template +void CheckToTensor(const std::vector& values, const std::shared_ptr typ, + const int32_t& element_size, const std::vector& element_shape, + const std::vector& element_permutation, + const std::vector& element_dim_names, + const std::vector& tensor_shape, + const std::vector& tensor_dim_names, + const std::vector& tensor_strides) { + auto buffer = Buffer::Wrap(values); + const std::shared_ptr element_type = fixed_size_list(typ, element_size); + std::vector> buffers = {nullptr, buffer}; + auto arr_data = std::make_shared(typ, values.size(), buffers); + auto arr = std::make_shared(arr_data); + ASSERT_OK_AND_ASSIGN(auto fsla_arr, FixedSizeListArray::FromArrays(arr, element_type)); + + ASSERT_OK_AND_ASSIGN( + auto expected_tensor, + Tensor::Make(typ, buffer, tensor_shape, tensor_strides, tensor_dim_names)); + const auto ext_type = + fixed_shape_tensor(typ, element_shape, element_permutation, element_dim_names); + + auto ext_arr = ExtensionType::WrapArray(ext_type, fsla_arr); + const auto tensor_array = std::static_pointer_cast(ext_arr); + ASSERT_OK_AND_ASSIGN(const auto actual_tensor, tensor_array->ToTensor()); + ASSERT_OK(actual_tensor->Validate()); + + ASSERT_EQ(actual_tensor->type(), expected_tensor->type()); + ASSERT_EQ(actual_tensor->shape(), expected_tensor->shape()); + ASSERT_EQ(actual_tensor->strides(), expected_tensor->strides()); + ASSERT_EQ(actual_tensor->dim_names(), expected_tensor->dim_names()); + ASSERT_TRUE(actual_tensor->data()->Equals(*expected_tensor->data())); + ASSERT_TRUE(actual_tensor->Equals(*expected_tensor)); +} + +TEST_F(TestExtensionType, ToTensor) { + std::vector float_values = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, + 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35}; + + auto element_sizes = std::vector{6, 6, 18, 18, 18, 18}; + + auto element_shapes = std::vector>{{2, 3}, {3, 2}, {3, 6}, + {6, 3}, {3, 2, 3}, {3, 2, 3}}; + auto tensor_shapes = std::vector>{ + {6, 2, 3}, {6, 2, 3}, {2, 3, 6}, {2, 3, 6}, {2, 3, 2, 3}, {2, 3, 2, 3}}; + + auto element_permutations = std::vector>{ + {0, 1}, {1, 0}, {0, 1}, {1, 0}, {0, 1, 2}, {2, 1, 0}}; + auto tensor_strides_32 = + std::vector>{{24, 12, 4}, {24, 4, 8}, {72, 24, 4}, + {72, 4, 12}, {72, 24, 12, 4}, {72, 4, 12, 24}}; + auto tensor_strides_64 = + std::vector>{{48, 24, 8}, {48, 8, 16}, {144, 48, 8}, + {144, 8, 24}, {144, 48, 24, 8}, {144, 8, 24, 48}}; + + auto element_dim_names = std::vector>{ + {"y", "z"}, {"z", "y"}, {"y", "z"}, {"z", "y"}, {"H", "W", "C"}, {"H", "W", "C"}}; + auto tensor_dim_names = std::vector>{ + {"", "y", "z"}, {"", "y", "z"}, {"", "y", "z"}, + {"", "y", "z"}, {"", "H", "W", "C"}, {"", "C", "W", "H"}}; + + for (size_t i = 0; i < element_shapes.size(); i++) { + CheckToTensor(float_values, float32(), element_sizes[i], element_shapes[i], + element_permutations[i], element_dim_names[i], tensor_shapes[i], + tensor_dim_names[i], tensor_strides_32[i]); + CheckToTensor(values_, int64(), element_sizes[i], element_shapes[i], + element_permutations[i], element_dim_names[i], tensor_shapes[i], + tensor_dim_names[i], tensor_strides_64[i]); + } +} + void CheckTensorRoundtrip(const std::shared_ptr& tensor) { ASSERT_OK_AND_ASSIGN(auto ext_arr, FixedShapeTensorArray::FromTensor(tensor)); ASSERT_OK_AND_ASSIGN(auto tensor_from_array, ext_arr->ToTensor()); @@ -364,7 +450,7 @@ TEST_F(TestExtensionType, SliceTensor) { Tensor::Make(value_type_, Buffer::Wrap(values_partial_), shape_partial_)); ASSERT_EQ(tensor->strides(), tensor_strides_); ASSERT_EQ(tensor_partial->strides(), tensor_strides_); - auto ext_type = fixed_shape_tensor(value_type_, cell_shape_, {}, dim_names_); + auto ext_type = fixed_shape_tensor(value_type_, element_shape_, {}, dim_names_); auto exact_ext_type = internal::checked_pointer_cast(ext_type_); ASSERT_OK_AND_ASSIGN(auto ext_arr, FixedShapeTensorArray::FromTensor(tensor)); @@ -404,11 +490,11 @@ TEST_F(TestExtensionType, ComputeStrides) { auto exact_ext_type = internal::checked_pointer_cast(ext_type_); auto ext_type_1 = internal::checked_pointer_cast( - fixed_shape_tensor(int64(), cell_shape_, {}, dim_names_)); + fixed_shape_tensor(int64(), element_shape_, {}, dim_names_)); auto ext_type_2 = internal::checked_pointer_cast( - fixed_shape_tensor(int64(), cell_shape_, {}, dim_names_)); + fixed_shape_tensor(int64(), element_shape_, {}, dim_names_)); auto ext_type_3 = internal::checked_pointer_cast( - fixed_shape_tensor(int32(), cell_shape_, {}, dim_names_)); + fixed_shape_tensor(int32(), element_shape_, {}, dim_names_)); ASSERT_TRUE(ext_type_1->Equals(*ext_type_2)); ASSERT_FALSE(ext_type_1->Equals(*ext_type_3)); @@ -462,4 +548,96 @@ TEST_F(TestExtensionType, ToString) { ASSERT_EQ(expected_3, result_3); } +TEST_F(TestExtensionType, GetTensor) { + auto arr = ArrayFromJSON(element_type_, + "[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]," + "[12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]]"); + auto element_values = + std::vector>{{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, + {12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}}; + + auto ext_type = fixed_shape_tensor(value_type_, element_shape_, {}, dim_names_); + auto permuted_ext_type = fixed_shape_tensor(value_type_, {3, 4}, {1, 0}, {"x", "y"}); + auto exact_ext_type = internal::checked_pointer_cast(ext_type); + auto exact_permuted_ext_type = + internal::checked_pointer_cast(permuted_ext_type); + + auto array = std::static_pointer_cast( + ExtensionType::WrapArray(ext_type, arr)); + auto permuted_array = std::static_pointer_cast( + ExtensionType::WrapArray(permuted_ext_type, arr)); + + for (size_t i = 0; i < element_values.size(); i++) { + // Get tensor from extension array with trivial permutation + ASSERT_OK_AND_ASSIGN(auto scalar, array->GetScalar(i)); + auto actual_ext_scalar = internal::checked_pointer_cast(scalar); + ASSERT_OK_AND_ASSIGN(auto actual_tensor, + exact_ext_type->MakeTensor(actual_ext_scalar)); + ASSERT_OK(actual_tensor->Validate()); + ASSERT_OK_AND_ASSIGN(auto expected_tensor, + Tensor::Make(value_type_, Buffer::Wrap(element_values[i]), + {3, 4}, {}, {"x", "y"})); + ASSERT_EQ(expected_tensor->shape(), actual_tensor->shape()); + ASSERT_EQ(expected_tensor->dim_names(), actual_tensor->dim_names()); + ASSERT_EQ(expected_tensor->strides(), actual_tensor->strides()); + ASSERT_EQ(actual_tensor->strides(), std::vector({32, 8})); + ASSERT_EQ(expected_tensor->type(), actual_tensor->type()); + ASSERT_TRUE(expected_tensor->Equals(*actual_tensor)); + + // Get tensor from extension array with non-trivial permutation + ASSERT_OK_AND_ASSIGN(auto expected_permuted_tensor, + Tensor::Make(value_type_, Buffer::Wrap(element_values[i]), + {4, 3}, {8, 24}, {"y", "x"})); + ASSERT_OK_AND_ASSIGN(scalar, permuted_array->GetScalar(i)); + ASSERT_OK_AND_ASSIGN(auto actual_permuted_tensor, + exact_permuted_ext_type->MakeTensor( + internal::checked_pointer_cast(scalar))); + ASSERT_OK(actual_permuted_tensor->Validate()); + ASSERT_EQ(expected_permuted_tensor->strides(), actual_permuted_tensor->strides()); + ASSERT_EQ(expected_permuted_tensor->shape(), actual_permuted_tensor->shape()); + ASSERT_EQ(expected_permuted_tensor->dim_names(), actual_permuted_tensor->dim_names()); + ASSERT_EQ(expected_permuted_tensor->type(), actual_permuted_tensor->type()); + ASSERT_EQ(expected_permuted_tensor->is_contiguous(), + actual_permuted_tensor->is_contiguous()); + ASSERT_EQ(expected_permuted_tensor->is_column_major(), + actual_permuted_tensor->is_column_major()); + ASSERT_TRUE(expected_permuted_tensor->Equals(*actual_permuted_tensor)); + } + + // Test null values fail + auto element_type = fixed_size_list(int64(), 1); + auto fsla_arr = ArrayFromJSON(element_type, "[[1], [null], null]"); + ext_type = fixed_shape_tensor(int64(), {1}); + exact_ext_type = internal::checked_pointer_cast(ext_type); + auto ext_arr = ExtensionType::WrapArray(ext_type, fsla_arr); + auto tensor_array = internal::checked_pointer_cast(ext_arr); + + ASSERT_OK_AND_ASSIGN(auto scalar, tensor_array->GetScalar(0)); + ASSERT_OK_AND_ASSIGN(auto tensor, + exact_ext_type->MakeTensor( + internal::checked_pointer_cast(scalar))); + + ASSERT_OK_AND_ASSIGN(scalar, tensor_array->GetScalar(1)); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, testing::HasSubstr("Invalid: Cannot convert data with nulls to Tensor."), + exact_ext_type->MakeTensor( + internal::checked_pointer_cast(scalar))); + + ASSERT_OK_AND_ASSIGN(scalar, tensor_array->GetScalar(2)); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, testing::HasSubstr("Invalid: Cannot convert data with nulls to Tensor."), + exact_ext_type->MakeTensor( + internal::checked_pointer_cast(scalar))); + + element_type = list(utf8()); + ext_type = fixed_shape_tensor(utf8(), {1}); + exact_ext_type = internal::checked_pointer_cast(ext_type); + scalar = std::make_shared(ArrayFromJSON(element_type, R"([["a", "b"]])")); + auto ext_scalar = std::make_shared(scalar, ext_type); + EXPECT_RAISES_WITH_MESSAGE_THAT( + TypeError, + testing::HasSubstr("Type error: Cannot convert non-fixed-width values to Tensor."), + exact_ext_type->MakeTensor(ext_scalar)); +} + } // namespace arrow diff --git a/cpp/src/arrow/extension/tensor_internal.h b/cpp/src/arrow/extension/tensor_internal.h new file mode 100644 index 0000000000000..069880cb17c85 --- /dev/null +++ b/cpp/src/arrow/extension/tensor_internal.h @@ -0,0 +1,45 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include + +#include "arrow/status.h" +#include "arrow/util/print.h" + +namespace arrow::internal { + +ARROW_EXPORT +Status IsPermutationValid(const std::vector& permutation) { + const auto size = static_cast(permutation.size()); + std::vector dim_seen(size, 0); + + for (const auto p : permutation) { + if (p < 0 || p >= size || dim_seen[p] != 0) { + return Status::Invalid( + "Permutation indices for ", size, + " dimensional tensors must be unique and within [0, ", size - 1, + "] range. Got: ", ::arrow::internal::PrintVector{permutation, ","}); + } + dim_seen[p] = 1; + } + return Status::OK(); +} + +} // namespace arrow::internal diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index 1416f5f4346d9..b9661a2bdb208 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -3529,7 +3529,7 @@ cdef class ExtensionArray(Array): return result -class FixedShapeTensorArray(ExtensionArray): +cdef class FixedShapeTensorArray(ExtensionArray): """ Concrete class for fixed shape tensor extension arrays. @@ -3570,17 +3570,48 @@ class FixedShapeTensorArray(ExtensionArray): def to_numpy_ndarray(self): """ - Convert fixed shape tensor extension array to a numpy array (with dim+1). + Convert fixed shape tensor extension array to a multi-dimensional numpy.ndarray. - Note: ``permutation`` should be trivial (``None`` or ``[0, 1, ..., len(shape)-1]``). + The resulting ndarray will have (ndim + 1) dimensions. + The size of the first dimension will be the length of the fixed shape tensor array + and the rest of the dimensions will match the permuted shape of the fixed + shape tensor. + + The conversion is zero-copy. + + Returns + ------- + numpy.ndarray + Ndarray representing tensors in the fixed shape tensor array concatenated + along the first dimension. """ - if self.type.permutation is None or self.type.permutation == list(range(len(self.type.shape))): - np_flat = np.asarray(self.storage.flatten()) - numpy_tensor = np_flat.reshape((len(self),) + tuple(self.type.shape)) - return numpy_tensor - else: - raise ValueError( - 'Only non-permuted tensors can be converted to numpy tensors.') + + return self.to_tensor().to_numpy() + + def to_tensor(self): + """ + Convert fixed shape tensor extension array to a pyarrow.Tensor. + + The resulting Tensor will have (ndim + 1) dimensions. + The size of the first dimension will be the length of the fixed shape tensor array + and the rest of the dimensions will match the permuted shape of the fixed + shape tensor. + + The conversion is zero-copy. + + Returns + ------- + pyarrow.Tensor + Tensor representing tensors in the fixed shape tensor array concatenated + along the first dimension. + """ + + cdef: + CFixedShapeTensorArray* ext_array = (self.ap) + CResult[shared_ptr[CTensor]] ctensor + with nogil: + ctensor = ext_array.ToTensor() + return pyarrow_wrap_tensor(GetResultValue(ctensor)) @staticmethod def from_numpy_ndarray(obj): @@ -3588,9 +3619,7 @@ class FixedShapeTensorArray(ExtensionArray): Convert numpy tensors (ndarrays) to a fixed shape tensor extension array. The first dimension of ndarray will become the length of the fixed shape tensor array. - - Numpy array needs to be C-contiguous in memory - (``obj.flags["C_CONTIGUOUS"]==True``). + If input array data is not contiguous a copy will be made. Parameters ---------- @@ -3624,17 +3653,25 @@ class FixedShapeTensorArray(ExtensionArray): ] ] """ - if not obj.flags["C_CONTIGUOUS"]: - raise ValueError('The data in the numpy array need to be in a single, ' - 'C-style contiguous segment.') + + if len(obj.shape) < 2: + raise ValueError( + "Cannot convert 1D array or scalar to fixed shape tensor array") + if np.prod(obj.shape) == 0: + raise ValueError("Expected a non-empty ndarray") + + permutation = (-np.array(obj.strides)).argsort(kind='stable') + if permutation[0] != 0: + raise ValueError('First stride needs to be largest to ensure that ' + 'individual tensor data is contiguous in memory.') arrow_type = from_numpy_dtype(obj.dtype) - shape = obj.shape[1:] - size = obj.size / obj.shape[0] + shape = np.take(obj.shape, permutation) + values = np.ravel(obj, order="K") return ExtensionArray.from_storage( - fixed_shape_tensor(arrow_type, shape), - FixedSizeListArray.from_arrays(np.ravel(obj, order='C'), size) + fixed_shape_tensor(arrow_type, shape[1:], permutation=permutation[1:] - 1), + FixedSizeListArray.from_arrays(values, shape[1:].prod()) ) diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 74e92594b04e5..1b33b462e66c5 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -2695,26 +2695,26 @@ cdef extern from "arrow/extension_type.h" namespace "arrow": shared_ptr[CArray] storage() -cdef extern from "arrow/extension/fixed_shape_tensor.h" namespace "arrow::extension": +cdef extern from "arrow/extension/fixed_shape_tensor.h" namespace "arrow::extension" nogil: cdef cppclass CFixedShapeTensorType \ " arrow::extension::FixedShapeTensorType"(CExtensionType): + CResult[shared_ptr[CTensor]] MakeTensor(const shared_ptr[CExtensionScalar]& scalar) const + @staticmethod CResult[shared_ptr[CDataType]] Make(const shared_ptr[CDataType]& value_type, const vector[int64_t]& shape, const vector[int64_t]& permutation, const vector[c_string]& dim_names) - CResult[shared_ptr[CDataType]] Deserialize(const shared_ptr[CDataType] storage_type, - const c_string& serialized_data) const - - c_string Serialize() const - const shared_ptr[CDataType] value_type() const vector[int64_t] shape() const vector[int64_t] permutation() const vector[c_string] dim_names() + cdef cppclass CFixedShapeTensorArray \ + " arrow::extension::FixedShapeTensorArray"(CExtensionArray): + const CResult[shared_ptr[CTensor]] ToTensor() const cdef extern from "arrow/util/compression.h" namespace "arrow" nogil: cdef enum CCompressionType" arrow::Compression::type": diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi index 9a66dc81226d4..f51f9cb804481 100644 --- a/python/pyarrow/scalar.pxi +++ b/python/pyarrow/scalar.pxi @@ -1027,6 +1027,48 @@ cdef class ExtensionScalar(Scalar): return pyarrow_wrap_scalar( sp_scalar) +cdef class FixedShapeTensorScalar(ExtensionScalar): + """ + Concrete class for fixed shape tensor extension scalar. + """ + + def to_numpy(self): + """ + Convert fixed shape tensor scalar to a numpy.ndarray. + + The resulting ndarray's shape matches the permuted shape of the + fixed shape tensor scalar. + The conversion is zero-copy. + + Returns + ------- + numpy.ndarray + """ + return self.to_tensor().to_numpy() + + def to_tensor(self): + """ + Convert fixed shape tensor extension scalar to a pyarrow.Tensor, using shape + and strides derived from corresponding FixedShapeTensorType. + + The conversion is zero-copy. + + Returns + ------- + pyarrow.Tensor + Tensor represented stored in FixedShapeTensorScalar. + """ + cdef: + CFixedShapeTensorType* c_type = static_pointer_cast[CFixedShapeTensorType, CDataType]( + self.wrapped.get().type).get() + shared_ptr[CExtensionScalar] scalar = static_pointer_cast[CExtensionScalar, CScalar](self.wrapped) + shared_ptr[CTensor] ctensor + + with nogil: + ctensor = GetResultValue(c_type.MakeTensor(scalar)) + return pyarrow_wrap_tensor(ctensor) + + cdef dict _scalar_classes = { _Type_BOOL: BooleanScalar, _Type_UINT8: UInt8Scalar, diff --git a/python/pyarrow/tests/test_extension_type.py b/python/pyarrow/tests/test_extension_type.py index a88e20eefe098..c081d7ae8303c 100644 --- a/python/pyarrow/tests/test_extension_type.py +++ b/python/pyarrow/tests/test_extension_type.py @@ -1318,39 +1318,120 @@ def test_tensor_type(): assert tensor_type.permutation is None -def test_tensor_class_methods(): - tensor_type = pa.fixed_shape_tensor(pa.float32(), [2, 3]) - storage = pa.array([[1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6]], - pa.list_(pa.float32(), 6)) +@pytest.mark.parametrize("value_type", (np.int8(), np.int64(), np.float32())) +def test_tensor_class_methods(value_type): + from numpy.lib.stride_tricks import as_strided + arrow_type = pa.from_numpy_dtype(value_type) + + tensor_type = pa.fixed_shape_tensor(arrow_type, [2, 3]) + storage = pa.array([[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]], + pa.list_(arrow_type, 6)) arr = pa.ExtensionArray.from_storage(tensor_type, storage) expected = np.array( - [[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]], dtype=np.float32) - result = arr.to_numpy_ndarray() + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], dtype=value_type) + np.testing.assert_array_equal(arr.to_tensor(), expected) + np.testing.assert_array_equal(arr.to_numpy_ndarray(), expected) + + expected = np.array([[[7, 8, 9], [10, 11, 12]]], dtype=value_type) + result = arr[1:].to_numpy_ndarray() np.testing.assert_array_equal(result, expected) - expected = np.array([[[1, 2, 3], [4, 5, 6]]], dtype=np.float32) - result = arr[:1].to_numpy_ndarray() + values = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] + flat_arr = np.array(values[0], dtype=value_type) + bw = value_type.itemsize + storage = pa.array(values, pa.list_(arrow_type, 12)) + + tensor_type = pa.fixed_shape_tensor(arrow_type, [2, 2, 3], permutation=[0, 1, 2]) + result = pa.ExtensionArray.from_storage(tensor_type, storage) + expected = np.array( + [[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]], dtype=value_type) + np.testing.assert_array_equal(result.to_numpy_ndarray(), expected) + + result = flat_arr.reshape(1, 2, 3, 2) + expected = np.array( + [[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]], dtype=value_type) np.testing.assert_array_equal(result, expected) - arr = np.array( - [[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]], - dtype=np.float32, order="C") + tensor_type = pa.fixed_shape_tensor(arrow_type, [2, 2, 3], permutation=[0, 2, 1]) + result = pa.ExtensionArray.from_storage(tensor_type, storage) + expected = as_strided(flat_arr, shape=(1, 2, 3, 2), + strides=(bw * 12, bw * 6, bw, bw * 3)) + np.testing.assert_array_equal(result.to_numpy_ndarray(), expected) + + tensor_type = pa.fixed_shape_tensor(arrow_type, [2, 2, 3], permutation=[2, 0, 1]) + result = pa.ExtensionArray.from_storage(tensor_type, storage) + expected = as_strided(flat_arr, shape=(1, 3, 2, 2), + strides=(bw * 12, bw, bw * 6, bw * 2)) + np.testing.assert_array_equal(result.to_numpy_ndarray(), expected) + + assert result.type.permutation == [2, 0, 1] + assert result.type.shape == [2, 2, 3] + assert result.to_tensor().shape == (1, 3, 2, 2) + assert result.to_tensor().strides == (12 * bw, 1 * bw, 6 * bw, 2 * bw) + + +@pytest.mark.parametrize("value_type", (np.int8(), np.int64(), np.float32())) +def test_tensor_array_from_numpy(value_type): + from numpy.lib.stride_tricks import as_strided + arrow_type = pa.from_numpy_dtype(value_type) + + arr = np.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + dtype=value_type, order="C") tensor_array_from_numpy = pa.FixedShapeTensorArray.from_numpy_ndarray(arr) assert isinstance(tensor_array_from_numpy.type, pa.FixedShapeTensorType) - assert tensor_array_from_numpy.type.value_type == pa.float32() + assert tensor_array_from_numpy.type.value_type == arrow_type assert tensor_array_from_numpy.type.shape == [2, 3] - arr = np.array( - [[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]], - dtype=np.float32, order="F") - with pytest.raises(ValueError, match="C-style contiguous segment"): + arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], + dtype=value_type, order="F") + with pytest.raises(ValueError, match="First stride needs to be largest"): pa.FixedShapeTensorArray.from_numpy_ndarray(arr) - tensor_type = pa.fixed_shape_tensor(pa.int8(), [2, 2, 3], permutation=[0, 2, 1]) - storage = pa.array([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6]], pa.list_(pa.int8(), 12)) - arr = pa.ExtensionArray.from_storage(tensor_type, storage) - with pytest.raises(ValueError, match="non-permuted tensors"): - arr.to_numpy_ndarray() + flat_arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=value_type) + bw = value_type.itemsize + + arr = flat_arr.reshape(1, 3, 4) + tensor_array_from_numpy = pa.FixedShapeTensorArray.from_numpy_ndarray(arr) + assert tensor_array_from_numpy.type.shape == [3, 4] + assert tensor_array_from_numpy.type.permutation == [0, 1] + assert tensor_array_from_numpy.to_tensor() == pa.Tensor.from_numpy(arr) + + arr = as_strided(flat_arr, shape=(1, 2, 3, 2), + strides=(bw * 12, bw * 6, bw, bw * 3)) + tensor_array_from_numpy = pa.FixedShapeTensorArray.from_numpy_ndarray(arr) + assert tensor_array_from_numpy.type.shape == [2, 2, 3] + assert tensor_array_from_numpy.type.permutation == [0, 2, 1] + assert tensor_array_from_numpy.to_tensor() == pa.Tensor.from_numpy(arr) + + arr = flat_arr.reshape(1, 2, 3, 2) + result = pa.FixedShapeTensorArray.from_numpy_ndarray(arr) + expected = np.array( + [[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]], dtype=value_type) + np.testing.assert_array_equal(result.to_numpy_ndarray(), expected) + + arr = np.array([[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]], dtype=value_type) + expected = arr[1:] + result = pa.FixedShapeTensorArray.from_numpy_ndarray(arr)[1:].to_numpy_ndarray() + np.testing.assert_array_equal(result, expected) + + arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=value_type) + with pytest.raises(ValueError, match="Cannot convert 1D array or scalar to fixed"): + pa.FixedShapeTensorArray.from_numpy_ndarray(arr) + + arr = np.array(1, dtype=value_type) + with pytest.raises(ValueError, match="Cannot convert 1D array or scalar to fixed"): + pa.FixedShapeTensorArray.from_numpy_ndarray(arr) + + arr = np.array([], dtype=value_type) + + with pytest.raises(ValueError, match="Cannot convert 1D array or scalar to fixed"): + pa.FixedShapeTensorArray.from_numpy_ndarray(arr.reshape((0))) + + with pytest.raises(ValueError, match="Expected a non-empty ndarray"): + pa.FixedShapeTensorArray.from_numpy_ndarray(arr.reshape((0, 3, 2))) + + with pytest.raises(ValueError, match="Expected a non-empty ndarray"): + pa.FixedShapeTensorArray.from_numpy_ndarray(arr.reshape((3, 0, 2))) @pytest.mark.parametrize("tensor_type", ( diff --git a/python/pyarrow/types.pxi b/python/pyarrow/types.pxi index b6dc53d633543..af5a261de348e 100644 --- a/python/pyarrow/types.pxi +++ b/python/pyarrow/types.pxi @@ -1658,20 +1658,6 @@ cdef class FixedShapeTensorType(BaseExtensionType): else: return None - def __arrow_ext_serialize__(self): - """ - Serialized representation of metadata to reconstruct the type object. - """ - return self.tensor_ext_type.Serialize() - - @classmethod - def __arrow_ext_deserialize__(self, storage_type, serialized): - """ - Return an FixedShapeTensor type instance from the storage type and serialized - metadata. - """ - return self.tensor_ext_type.Deserialize(storage_type, serialized) - def __arrow_ext_class__(self): return FixedShapeTensorArray @@ -1679,6 +1665,9 @@ cdef class FixedShapeTensorType(BaseExtensionType): return fixed_shape_tensor, (self.value_type, self.shape, self.dim_names, self.permutation) + def __arrow_ext_scalar_class__(self): + return FixedShapeTensorScalar + _py_extension_type_auto_load = False @@ -4946,8 +4935,9 @@ def fixed_shape_tensor(DataType value_type, shape, dim_names=None, permutation=N cdef FixedShapeTensorType out = FixedShapeTensorType.__new__(FixedShapeTensorType) - c_tensor_ext_type = GetResultValue(CFixedShapeTensorType.Make( - value_type.sp_type, c_shape, c_permutation, c_dim_names)) + with nogil: + c_tensor_ext_type = GetResultValue(CFixedShapeTensorType.Make( + value_type.sp_type, c_shape, c_permutation, c_dim_names)) out.init(c_tensor_ext_type)