tests/cpp/test_tensor.cc - tvm-ffi - Git at Google

 /*
  * 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.
  */
 #include <gtest/gtest.h>
 #include <tvm/ffi/container/tensor.h>

 namespace {

 using namespace tvm::ffi;

 struct CPUNDAlloc {
   void AllocData(DLTensor* tensor) { tensor->data = malloc(GetDataSize(*tensor)); }
   void FreeData(DLTensor* tensor) { free(tensor->data); }
 };

 inline Tensor Empty(const Shape& shape, DLDataType dtype, DLDevice device) {
   return Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device);
 }

 inline Tensor EmptyStrided(const Shape& shape, const Shape& strides, DLDataType dtype,
                            DLDevice device) {
   return Tensor::FromNDAllocStrided(CPUNDAlloc(), shape, strides, dtype, device);
 }

 int TestEnvTensorAllocator(DLTensor* prototype, TVMFFIObjectHandle* out) {
   Shape shape(prototype->shape, prototype->shape + prototype->ndim);
   Tensor nd = Empty(shape, prototype->dtype, prototype->device);
   *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(nd));
   return 0;
 }

 int TestEnvTensorAllocatorError(DLTensor* prototype, TVMFFIObjectHandle* out) {
   TVMFFIErrorSetRaisedFromCStr("RuntimeError", "TestEnvTensorAllocatorError");
   return -1;
 }

 TEST(Tensor, Basic) {
   Tensor nd = Empty({1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
   Shape shape = nd.shape();
   Shape strides = nd.strides();
   EXPECT_EQ(shape.size(), 3);
   EXPECT_EQ(shape[0], 1);
   EXPECT_EQ(shape[1], 2);
   EXPECT_EQ(shape[2], 3);
   EXPECT_EQ(strides.size(), 3);
   EXPECT_EQ(strides[0], 6);
   EXPECT_EQ(strides[1], 3);
   EXPECT_EQ(strides[2], 1);
   EXPECT_EQ(nd.dtype(), DLDataType({kDLFloat, 32, 1}));
   for (int64_t i = 0; i < shape.Product(); ++i) {
     reinterpret_cast<float*>(nd.data_ptr())[i] = static_cast<float>(i);
   }

   EXPECT_EQ(nd.numel(), 6);
   EXPECT_EQ(nd.ndim(), 3);
   EXPECT_EQ(nd.data_ptr(), nd.GetDLTensorPtr()->data);

   Any any0 = nd;
   Tensor nd2 = any0.as<Tensor>().value();  // NOLINT(bugprone-unchecked-optional-access)
   EXPECT_EQ(nd2.dtype(), DLDataType({kDLFloat, 32, 1}));
   for (int64_t i = 0; i < shape.Product(); ++i) {
     EXPECT_EQ(reinterpret_cast<float*>(nd2.data_ptr())[i], i);
   }

   EXPECT_EQ(nd.IsContiguous(), true);
   EXPECT_EQ(nd2.use_count(), 3);

   Tensor nd3 = EmptyStrided({2, 3}, {1, 2}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
   Shape shape3 = nd3.shape();
   Shape strides3 = nd3.strides();
   EXPECT_EQ(shape3.size(), 2);
   EXPECT_EQ(shape3[0], 2);
   EXPECT_EQ(shape3[1], 3);
   EXPECT_EQ(strides3.size(), 2);
   EXPECT_EQ(strides3[0], 1);
   EXPECT_EQ(strides3[1], 2);
 }

 TEST(Tensor, EmptyTensorIsContiguous) {
   // An empty tensor (any shape dim == 0) is trivially contiguous regardless of
   // stride values.  This matches NumPy / PyTorch semantics.
   // Use strides that would normally fail the contiguity check to verify the
   // early-return path in IsContiguous().
   Tensor nd =
       EmptyStrided({4, 0, 4}, {0, 0, 0}, DLDataType({kDLInt, 16, 1}), DLDevice({kDLCPU, 0}));
   EXPECT_EQ(nd.numel(), 0);
   EXPECT_EQ(nd.IsContiguous(), true);
   EXPECT_EQ(nd.is_contiguous(), true);
 }

 TEST(Tensor, DLPack) {
   Tensor tensor = Empty({1, 2, 3}, DLDataType({kDLInt, 16, 1}), DLDevice({kDLCPU, 0}));
   DLManagedTensor* dlpack = tensor.ToDLPack();
   EXPECT_EQ(dlpack->dl_tensor.ndim, 3);
   EXPECT_EQ(dlpack->dl_tensor.shape[0], 1);
   EXPECT_EQ(dlpack->dl_tensor.shape[1], 2);
   EXPECT_EQ(dlpack->dl_tensor.shape[2], 3);
   EXPECT_EQ(dlpack->dl_tensor.dtype.code, kDLInt);
   EXPECT_EQ(dlpack->dl_tensor.dtype.bits, 16);
   EXPECT_EQ(dlpack->dl_tensor.dtype.lanes, 1);
   EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU);
   EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0);
   EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0);
   EXPECT_EQ(dlpack->dl_tensor.strides[0], 6);
   EXPECT_EQ(dlpack->dl_tensor.strides[1], 3);
   EXPECT_EQ(dlpack->dl_tensor.strides[2], 1);
   EXPECT_EQ(tensor.use_count(), 2);
   {
     Tensor tensor2 = Tensor::FromDLPack(dlpack);
     EXPECT_EQ(tensor2.use_count(), 1);
     EXPECT_EQ(tensor2.data_ptr(), tensor.data_ptr());
     EXPECT_EQ(tensor.use_count(), 2);
     EXPECT_EQ(tensor2.use_count(), 1);
   }
   EXPECT_EQ(tensor.use_count(), 1);
 }

 TEST(Tensor, DLPackVersioned) {
   DLDataType dtype = DLDataType({kDLFloat4_e2m1fn, 4, 1});
   EXPECT_EQ(GetDataSize(2, dtype), 2 * 4 / 8);
   Tensor tensor = Empty({2}, dtype, DLDevice({kDLCPU, 0}));
   DLManagedTensorVersioned* dlpack = tensor.ToDLPackVersioned();
   EXPECT_EQ(dlpack->version.major, DLPACK_MAJOR_VERSION);
   EXPECT_EQ(dlpack->version.minor, DLPACK_MINOR_VERSION);
   EXPECT_EQ(dlpack->dl_tensor.ndim, 1);
   EXPECT_EQ(dlpack->dl_tensor.shape[0], 2);
   EXPECT_EQ(dlpack->dl_tensor.dtype.code, kDLFloat4_e2m1fn);
   EXPECT_EQ(dlpack->dl_tensor.dtype.bits, 4);
   EXPECT_EQ(dlpack->dl_tensor.dtype.lanes, 1);
   EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU);
   EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0);
   EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0);
   EXPECT_EQ(dlpack->dl_tensor.strides[0], 1);

   EXPECT_EQ(tensor.use_count(), 2);
   {
     Tensor tensor2 = Tensor::FromDLPackVersioned(dlpack);
     EXPECT_EQ(tensor2.use_count(), 1);
     EXPECT_EQ(tensor2.data_ptr(), tensor.data_ptr());
     EXPECT_EQ(tensor.use_count(), 2);
     EXPECT_EQ(tensor2.use_count(), 1);
   }
   EXPECT_EQ(tensor.use_count(), 1);
 }

 TEST(Tensor, EnvAlloc) {
   // Test successful allocation
   Tensor tensor = Tensor::FromEnvAlloc(TestEnvTensorAllocator, {1, 2, 3},
                                        DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
   EXPECT_EQ(tensor.use_count(), 1);
   EXPECT_EQ(tensor.shape().size(), 3);
   EXPECT_EQ(tensor.size(0), 1);
   EXPECT_EQ(tensor.size(1), 2);
   EXPECT_EQ(tensor.size(2), 3);
   EXPECT_EQ(tensor.size(-3), 1);
   EXPECT_EQ(tensor.size(-2), 2);
   EXPECT_EQ(tensor.size(-1), 3);
   EXPECT_EQ(tensor.stride(0), 6);
   EXPECT_EQ(tensor.stride(1), 3);
   EXPECT_EQ(tensor.stride(2), 1);
   EXPECT_EQ(tensor.stride(-3), 6);
   EXPECT_EQ(tensor.stride(-2), 3);
   EXPECT_EQ(tensor.stride(-1), 1);
   EXPECT_EQ(tensor.dtype().code, kDLFloat);
   EXPECT_EQ(tensor.dtype().bits, 32);
   EXPECT_EQ(tensor.dtype().lanes, 1);
   EXPECT_EQ(tensor.device().device_type, kDLCPU);
   EXPECT_EQ(tensor.device().device_id, 0);
   EXPECT_NE(tensor.data_ptr(), nullptr);
 }

 TEST(Tensor, EnvAllocError) {
   // Test error handling in DLPackAlloc
   EXPECT_THROW(
       {
         Tensor::FromEnvAlloc(TestEnvTensorAllocatorError, {1, 2, 3}, DLDataType({kDLFloat, 32, 1}),
                              DLDevice({kDLCPU, 0}));
       },
       tvm::ffi::Error);
 }

 TEST(Tensor, TensorView) {
   Tensor tensor = Empty({1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
   TensorView tensor_view = tensor;

   EXPECT_EQ(tensor_view.shape().size(), 3);
   EXPECT_EQ(tensor_view.shape()[0], 1);
   EXPECT_EQ(tensor_view.shape()[1], 2);
   EXPECT_EQ(tensor_view.shape()[2], 3);
   EXPECT_EQ(tensor_view.dtype().code, kDLFloat);
   EXPECT_EQ(tensor_view.dtype().bits, 32);
   EXPECT_EQ(tensor_view.dtype().lanes, 1);

   AnyView result = tensor_view;
   EXPECT_EQ(result.type_index(), TypeIndex::kTVMFFIDLTensorPtr);
   // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
   TensorView tensor_view2 = result.as<TensorView>().value();
   EXPECT_EQ(tensor_view2.shape().size(), 3);
   EXPECT_EQ(tensor_view2.shape()[0], 1);
   EXPECT_EQ(tensor_view2.shape()[1], 2);
   EXPECT_EQ(tensor_view2.shape()[2], 3);
   EXPECT_EQ(tensor_view2.dtype().code, kDLFloat);
   EXPECT_EQ(tensor_view2.dtype().bits, 32);
   EXPECT_EQ(tensor_view2.dtype().lanes, 1);
 }

 TEST(Tensor, TensorViewAsStrided) {
   // Create a base tensor with shape [2, 3] = 6 elements
   Tensor tensor = Empty({2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));

   // Fill with sequential values: [0, 1, 2, 3, 4, 5]
   float* data = reinterpret_cast<float*>(tensor.data_ptr());
   size_t element_capacity = GetDataSize(tensor) / sizeof(float);
   ASSERT_EQ(element_capacity, static_cast<size_t>(tensor.numel()));
   for (size_t i = 0; i < element_capacity; ++i) {
     data[i] = static_cast<float>(i);
   }

   TensorView tensor_view = tensor;
   void* original_data_ptr = tensor_view.data_ptr();
   EXPECT_EQ(tensor_view.byte_offset(), 0);

   // Create a strided view with shape [3, 2] and custom strides
   // Use local variables to ensure they stay in scope for the TensorView
   Shape new_shape = {3, 2};
   Shape new_strides = {1, 3};
   TensorView strided_view = tensor_view.as_strided(new_shape, new_strides);

   // Verify the view has correct shape and strides
   EXPECT_EQ(strided_view.shape().size(), 2);
   EXPECT_EQ(strided_view.shape()[0], 3);
   EXPECT_EQ(strided_view.shape()[1], 2);
   EXPECT_EQ(strided_view.strides().size(), 2);
   EXPECT_EQ(strided_view.strides()[0], 1);
   EXPECT_EQ(strided_view.strides()[1], 3);

   // Verify the view shares the same underlying data pointer (no offset)
   EXPECT_EQ(strided_view.data_ptr(), original_data_ptr);
   EXPECT_EQ(strided_view.byte_offset(), 0);
   EXPECT_EQ(strided_view.dtype(), tensor_view.dtype());

   // Test with element_offset - for float32, 1 element = 4 bytes
   Shape offset_shape = {2, 2};
   Shape offset_strides = {3, 1};
   int64_t element_offset = 1;
   TensorView offset_view = tensor_view.as_strided(offset_shape, offset_strides, element_offset);

   EXPECT_EQ(offset_view.shape().size(), 2);
   EXPECT_EQ(offset_view.shape()[0], 2);
   EXPECT_EQ(offset_view.shape()[1], 2);
   EXPECT_EQ(offset_view.strides().size(), 2);
   EXPECT_EQ(offset_view.strides()[0], 3);
   EXPECT_EQ(offset_view.strides()[1], 1);

   // For CPU (direct address device), byte_offset should be added to data pointer
   // and byte_offset field should be 0
   // element_offset=1 for float32 = 4 bytes
   size_t expected_byte_offset =
       GetDataSize(static_cast<size_t>(element_offset), DLDataType({kDLFloat, 32, 1}));
   EXPECT_EQ(expected_byte_offset, 4);  // 1 element * 32 bits / 8 = 4 bytes

   // The data pointer should be advanced by 4 bytes (1 float element)
   void* expected_offset_ptr = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset;
   EXPECT_EQ(offset_view.data_ptr(), expected_offset_ptr);
   EXPECT_EQ(offset_view.byte_offset(), 0);  // Should be 0 for direct address devices

   // Verify data access through the offset view
   float* offset_data = reinterpret_cast<float*>(offset_view.data_ptr());
   EXPECT_EQ(offset_data[0 * 3 + 0 * 1], 1.0f);  // Points to data[1]
   EXPECT_EQ(offset_data[1 * 3 + 0 * 1], 4.0f);  // Points to data[4]

   // Test with larger element_offset
   int64_t element_offset2 = 2;
   Shape offset_shape2 = {1, 2};
   Shape offset_strides2 = {3, 1};
   TensorView offset_view2 = tensor_view.as_strided(offset_shape2, offset_strides2, element_offset2);
   size_t expected_byte_offset2 =
       GetDataSize(static_cast<size_t>(element_offset2), DLDataType({kDLFloat, 32, 1}));
   EXPECT_EQ(expected_byte_offset2, 8);  // 2 elements * 32 bits / 8 = 8 bytes
   void* expected_offset_ptr2 = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset2;
   EXPECT_EQ(offset_view2.data_ptr(), expected_offset_ptr2);
   EXPECT_EQ(offset_view2.byte_offset(), 0);

   float* offset_data2 = reinterpret_cast<float*>(offset_view2.data_ptr());
   EXPECT_EQ(offset_data2[0 * 3 + 0 * 1], 2.0f);  // Points to data[2]
 }

 TEST(Tensor, AsStrided) {
   // Create a base tensor with shape [2, 3] = 6 elements
   Tensor tensor = Empty({2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));

   // Fill with sequential values: [0, 1, 2, 3, 4, 5]
   float* data = reinterpret_cast<float*>(tensor.data_ptr());
   size_t element_capacity = GetDataSize(tensor) / sizeof(float);
   ASSERT_EQ(element_capacity, static_cast<size_t>(tensor.numel()));
   for (size_t i = 0; i < element_capacity; ++i) {
     data[i] = static_cast<float>(i);
   }

   void* original_data_ptr = tensor.data_ptr();
   EXPECT_EQ(tensor.byte_offset(), 0);

   // Create a strided view with shape [3, 2] and custom strides
   Shape new_shape = {3, 2};
   Shape new_strides = {1, 3};
   Tensor strided_view = tensor.as_strided(new_shape, new_strides);

   // Verify the view has correct shape and strides
   EXPECT_EQ(strided_view.shape().size(), 2);
   EXPECT_EQ(strided_view.shape()[0], 3);
   EXPECT_EQ(strided_view.shape()[1], 2);
   EXPECT_EQ(strided_view.strides().size(), 2);
   EXPECT_EQ(strided_view.strides()[0], 1);
   EXPECT_EQ(strided_view.strides()[1], 3);

   // Verify the view shares the same underlying data pointer (no offset)
   EXPECT_EQ(strided_view.data_ptr(), original_data_ptr);
   EXPECT_EQ(strided_view.byte_offset(), 0);
   EXPECT_EQ(strided_view.dtype(), tensor.dtype());

   // Test with element_offset - for float32, 1 element = 4 bytes
   Shape offset_shape = {2, 2};
   Shape offset_strides = {3, 1};
   int64_t element_offset = 1;
   Tensor offset_view = tensor.as_strided(offset_shape, offset_strides, element_offset);

   EXPECT_EQ(offset_view.shape().size(), 2);
   EXPECT_EQ(offset_view.shape()[0], 2);
   EXPECT_EQ(offset_view.shape()[1], 2);
   EXPECT_EQ(offset_view.strides().size(), 2);
   EXPECT_EQ(offset_view.strides()[0], 3);
   EXPECT_EQ(offset_view.strides()[1], 1);

   // For CPU (direct address device), byte_offset should be added to data pointer
   // and byte_offset field should be 0
   // element_offset=1 for float32 = 4 bytes
   size_t expected_byte_offset =
       GetDataSize(static_cast<size_t>(element_offset), DLDataType({kDLFloat, 32, 1}));
   EXPECT_EQ(expected_byte_offset, 4);  // 1 element * 32 bits / 8 = 4 bytes

   // The data pointer should be advanced by 4 bytes (1 float element)
   void* expected_offset_ptr = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset;
   EXPECT_EQ(offset_view.data_ptr(), expected_offset_ptr);
   EXPECT_EQ(offset_view.byte_offset(), 0);  // Should be 0 for direct address devices

   // Verify data access through the offset view
   float* offset_data = reinterpret_cast<float*>(offset_view.data_ptr());
   EXPECT_EQ(offset_data[0 * 3 + 0 * 1], 1.0f);  // Points to data[1]
   EXPECT_EQ(offset_data[1 * 3 + 0 * 1], 4.0f);  // Points to data[4]

   // Test with larger element_offset
   int64_t element_offset2 = 2;
   Tensor offset_view2 = tensor.as_strided({1, 2}, {3, 1}, element_offset2);
   size_t expected_byte_offset2 =
       GetDataSize(static_cast<size_t>(element_offset2), DLDataType({kDLFloat, 32, 1}));
   EXPECT_EQ(expected_byte_offset2, 8);  // 2 elements * 32 bits / 8 = 8 bytes
   void* expected_offset_ptr2 = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset2;
   EXPECT_EQ(offset_view2.data_ptr(), expected_offset_ptr2);
   EXPECT_EQ(offset_view2.byte_offset(), 0);

   float* offset_data2 = reinterpret_cast<float*>(offset_view2.data_ptr());
   EXPECT_EQ(offset_data2[0 * 3 + 0 * 1], 2.0f);  // Points to data[2]
 }

 TEST(Tensor, SizeStrideOutOfBounds) {
   Tensor tensor = Empty({2, 3, 4}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
   EXPECT_THROW({ tensor.size(3); }, tvm::ffi::Error);
   EXPECT_THROW({ tensor.size(-4); }, tvm::ffi::Error);
   EXPECT_THROW({ tensor.stride(3); }, tvm::ffi::Error);
   EXPECT_THROW({ tensor.stride(-4); }, tvm::ffi::Error);

   TensorView tensor_view = tensor;
   EXPECT_THROW({ tensor_view.size(3); }, tvm::ffi::Error);
   EXPECT_THROW({ tensor_view.size(-4); }, tvm::ffi::Error);
   EXPECT_THROW({ tensor_view.stride(3); }, tvm::ffi::Error);
   EXPECT_THROW({ tensor_view.stride(-4); }, tvm::ffi::Error);
 }

 }  // namespace
	/*
	* 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.
	*/
	#include <gtest/gtest.h>
	#include <tvm/ffi/container/tensor.h>

	namespace {

	using namespace tvm::ffi;

	struct CPUNDAlloc {
	void AllocData(DLTensor* tensor) { tensor->data = malloc(GetDataSize(*tensor)); }
	void FreeData(DLTensor* tensor) { free(tensor->data); }
	};

	inline Tensor Empty(const Shape& shape, DLDataType dtype, DLDevice device) {
	return Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device);
	}

	inline Tensor EmptyStrided(const Shape& shape, const Shape& strides, DLDataType dtype,
	DLDevice device) {
	return Tensor::FromNDAllocStrided(CPUNDAlloc(), shape, strides, dtype, device);
	}

	int TestEnvTensorAllocator(DLTensor* prototype, TVMFFIObjectHandle* out) {
	Shape shape(prototype->shape, prototype->shape + prototype->ndim);
	Tensor nd = Empty(shape, prototype->dtype, prototype->device);
	*out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(nd));
	return 0;
	}

	int TestEnvTensorAllocatorError(DLTensor* prototype, TVMFFIObjectHandle* out) {
	TVMFFIErrorSetRaisedFromCStr("RuntimeError", "TestEnvTensorAllocatorError");
	return -1;
	}

	TEST(Tensor, Basic) {
	Tensor nd = Empty({1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
	Shape shape = nd.shape();
	Shape strides = nd.strides();
	EXPECT_EQ(shape.size(), 3);
	EXPECT_EQ(shape[0], 1);
	EXPECT_EQ(shape[1], 2);
	EXPECT_EQ(shape[2], 3);
	EXPECT_EQ(strides.size(), 3);
	EXPECT_EQ(strides[0], 6);
	EXPECT_EQ(strides[1], 3);
	EXPECT_EQ(strides[2], 1);
	EXPECT_EQ(nd.dtype(), DLDataType({kDLFloat, 32, 1}));
	for (int64_t i = 0; i < shape.Product(); ++i) {
	reinterpret_cast<float*>(nd.data_ptr())[i] = static_cast<float>(i);
	}

	EXPECT_EQ(nd.numel(), 6);
	EXPECT_EQ(nd.ndim(), 3);
	EXPECT_EQ(nd.data_ptr(), nd.GetDLTensorPtr()->data);

	Any any0 = nd;
	Tensor nd2 = any0.as<Tensor>().value(); // NOLINT(bugprone-unchecked-optional-access)
	EXPECT_EQ(nd2.dtype(), DLDataType({kDLFloat, 32, 1}));
	for (int64_t i = 0; i < shape.Product(); ++i) {
	EXPECT_EQ(reinterpret_cast<float*>(nd2.data_ptr())[i], i);
	}

	EXPECT_EQ(nd.IsContiguous(), true);
	EXPECT_EQ(nd2.use_count(), 3);

	Tensor nd3 = EmptyStrided({2, 3}, {1, 2}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
	Shape shape3 = nd3.shape();
	Shape strides3 = nd3.strides();
	EXPECT_EQ(shape3.size(), 2);
	EXPECT_EQ(shape3[0], 2);
	EXPECT_EQ(shape3[1], 3);
	EXPECT_EQ(strides3.size(), 2);
	EXPECT_EQ(strides3[0], 1);
	EXPECT_EQ(strides3[1], 2);
	}

	TEST(Tensor, EmptyTensorIsContiguous) {
	// An empty tensor (any shape dim == 0) is trivially contiguous regardless of
	// stride values. This matches NumPy / PyTorch semantics.
	// Use strides that would normally fail the contiguity check to verify the
	// early-return path in IsContiguous().
	Tensor nd =
	EmptyStrided({4, 0, 4}, {0, 0, 0}, DLDataType({kDLInt, 16, 1}), DLDevice({kDLCPU, 0}));
	EXPECT_EQ(nd.numel(), 0);
	EXPECT_EQ(nd.IsContiguous(), true);
	EXPECT_EQ(nd.is_contiguous(), true);
	}

	TEST(Tensor, DLPack) {
	Tensor tensor = Empty({1, 2, 3}, DLDataType({kDLInt, 16, 1}), DLDevice({kDLCPU, 0}));
	DLManagedTensor* dlpack = tensor.ToDLPack();
	EXPECT_EQ(dlpack->dl_tensor.ndim, 3);
	EXPECT_EQ(dlpack->dl_tensor.shape[0], 1);
	EXPECT_EQ(dlpack->dl_tensor.shape[1], 2);
	EXPECT_EQ(dlpack->dl_tensor.shape[2], 3);
	EXPECT_EQ(dlpack->dl_tensor.dtype.code, kDLInt);
	EXPECT_EQ(dlpack->dl_tensor.dtype.bits, 16);
	EXPECT_EQ(dlpack->dl_tensor.dtype.lanes, 1);
	EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU);
	EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0);
	EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0);
	EXPECT_EQ(dlpack->dl_tensor.strides[0], 6);
	EXPECT_EQ(dlpack->dl_tensor.strides[1], 3);
	EXPECT_EQ(dlpack->dl_tensor.strides[2], 1);
	EXPECT_EQ(tensor.use_count(), 2);
	{
	Tensor tensor2 = Tensor::FromDLPack(dlpack);
	EXPECT_EQ(tensor2.use_count(), 1);
	EXPECT_EQ(tensor2.data_ptr(), tensor.data_ptr());
	EXPECT_EQ(tensor.use_count(), 2);
	EXPECT_EQ(tensor2.use_count(), 1);
	}
	EXPECT_EQ(tensor.use_count(), 1);
	}

	TEST(Tensor, DLPackVersioned) {
	DLDataType dtype = DLDataType({kDLFloat4_e2m1fn, 4, 1});
	EXPECT_EQ(GetDataSize(2, dtype), 2 * 4 / 8);
	Tensor tensor = Empty({2}, dtype, DLDevice({kDLCPU, 0}));
	DLManagedTensorVersioned* dlpack = tensor.ToDLPackVersioned();
	EXPECT_EQ(dlpack->version.major, DLPACK_MAJOR_VERSION);
	EXPECT_EQ(dlpack->version.minor, DLPACK_MINOR_VERSION);
	EXPECT_EQ(dlpack->dl_tensor.ndim, 1);
	EXPECT_EQ(dlpack->dl_tensor.shape[0], 2);
	EXPECT_EQ(dlpack->dl_tensor.dtype.code, kDLFloat4_e2m1fn);
	EXPECT_EQ(dlpack->dl_tensor.dtype.bits, 4);
	EXPECT_EQ(dlpack->dl_tensor.dtype.lanes, 1);
	EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU);
	EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0);
	EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0);
	EXPECT_EQ(dlpack->dl_tensor.strides[0], 1);

	EXPECT_EQ(tensor.use_count(), 2);
	{
	Tensor tensor2 = Tensor::FromDLPackVersioned(dlpack);
	EXPECT_EQ(tensor2.use_count(), 1);
	EXPECT_EQ(tensor2.data_ptr(), tensor.data_ptr());
	EXPECT_EQ(tensor.use_count(), 2);
	EXPECT_EQ(tensor2.use_count(), 1);
	}
	EXPECT_EQ(tensor.use_count(), 1);
	}

	TEST(Tensor, EnvAlloc) {
	// Test successful allocation
	Tensor tensor = Tensor::FromEnvAlloc(TestEnvTensorAllocator, {1, 2, 3},
	DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
	EXPECT_EQ(tensor.use_count(), 1);
	EXPECT_EQ(tensor.shape().size(), 3);
	EXPECT_EQ(tensor.size(0), 1);
	EXPECT_EQ(tensor.size(1), 2);
	EXPECT_EQ(tensor.size(2), 3);
	EXPECT_EQ(tensor.size(-3), 1);
	EXPECT_EQ(tensor.size(-2), 2);
	EXPECT_EQ(tensor.size(-1), 3);
	EXPECT_EQ(tensor.stride(0), 6);
	EXPECT_EQ(tensor.stride(1), 3);
	EXPECT_EQ(tensor.stride(2), 1);
	EXPECT_EQ(tensor.stride(-3), 6);
	EXPECT_EQ(tensor.stride(-2), 3);
	EXPECT_EQ(tensor.stride(-1), 1);
	EXPECT_EQ(tensor.dtype().code, kDLFloat);
	EXPECT_EQ(tensor.dtype().bits, 32);
	EXPECT_EQ(tensor.dtype().lanes, 1);
	EXPECT_EQ(tensor.device().device_type, kDLCPU);
	EXPECT_EQ(tensor.device().device_id, 0);
	EXPECT_NE(tensor.data_ptr(), nullptr);
	}

	TEST(Tensor, EnvAllocError) {
	// Test error handling in DLPackAlloc
	EXPECT_THROW(
	{
	Tensor::FromEnvAlloc(TestEnvTensorAllocatorError, {1, 2, 3}, DLDataType({kDLFloat, 32, 1}),
	DLDevice({kDLCPU, 0}));
	},
	tvm::ffi::Error);
	}

	TEST(Tensor, TensorView) {
	Tensor tensor = Empty({1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
	TensorView tensor_view = tensor;

	EXPECT_EQ(tensor_view.shape().size(), 3);
	EXPECT_EQ(tensor_view.shape()[0], 1);
	EXPECT_EQ(tensor_view.shape()[1], 2);
	EXPECT_EQ(tensor_view.shape()[2], 3);
	EXPECT_EQ(tensor_view.dtype().code, kDLFloat);
	EXPECT_EQ(tensor_view.dtype().bits, 32);
	EXPECT_EQ(tensor_view.dtype().lanes, 1);

	AnyView result = tensor_view;
	EXPECT_EQ(result.type_index(), TypeIndex::kTVMFFIDLTensorPtr);
	// NOLINTNEXTLINE(bugprone-unchecked-optional-access)
	TensorView tensor_view2 = result.as<TensorView>().value();
	EXPECT_EQ(tensor_view2.shape().size(), 3);
	EXPECT_EQ(tensor_view2.shape()[0], 1);
	EXPECT_EQ(tensor_view2.shape()[1], 2);
	EXPECT_EQ(tensor_view2.shape()[2], 3);
	EXPECT_EQ(tensor_view2.dtype().code, kDLFloat);
	EXPECT_EQ(tensor_view2.dtype().bits, 32);
	EXPECT_EQ(tensor_view2.dtype().lanes, 1);
	}

	TEST(Tensor, TensorViewAsStrided) {
	// Create a base tensor with shape [2, 3] = 6 elements
	Tensor tensor = Empty({2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));

	// Fill with sequential values: [0, 1, 2, 3, 4, 5]
	float* data = reinterpret_cast<float*>(tensor.data_ptr());
	size_t element_capacity = GetDataSize(tensor) / sizeof(float);
	ASSERT_EQ(element_capacity, static_cast<size_t>(tensor.numel()));
	for (size_t i = 0; i < element_capacity; ++i) {
	data[i] = static_cast<float>(i);
	}

	TensorView tensor_view = tensor;
	void* original_data_ptr = tensor_view.data_ptr();
	EXPECT_EQ(tensor_view.byte_offset(), 0);

	// Create a strided view with shape [3, 2] and custom strides
	// Use local variables to ensure they stay in scope for the TensorView
	Shape new_shape = {3, 2};
	Shape new_strides = {1, 3};
	TensorView strided_view = tensor_view.as_strided(new_shape, new_strides);

	// Verify the view has correct shape and strides
	EXPECT_EQ(strided_view.shape().size(), 2);
	EXPECT_EQ(strided_view.shape()[0], 3);
	EXPECT_EQ(strided_view.shape()[1], 2);
	EXPECT_EQ(strided_view.strides().size(), 2);
	EXPECT_EQ(strided_view.strides()[0], 1);
	EXPECT_EQ(strided_view.strides()[1], 3);

	// Verify the view shares the same underlying data pointer (no offset)
	EXPECT_EQ(strided_view.data_ptr(), original_data_ptr);
	EXPECT_EQ(strided_view.byte_offset(), 0);
	EXPECT_EQ(strided_view.dtype(), tensor_view.dtype());

	// Test with element_offset - for float32, 1 element = 4 bytes
	Shape offset_shape = {2, 2};
	Shape offset_strides = {3, 1};
	int64_t element_offset = 1;
	TensorView offset_view = tensor_view.as_strided(offset_shape, offset_strides, element_offset);

	EXPECT_EQ(offset_view.shape().size(), 2);
	EXPECT_EQ(offset_view.shape()[0], 2);
	EXPECT_EQ(offset_view.shape()[1], 2);
	EXPECT_EQ(offset_view.strides().size(), 2);
	EXPECT_EQ(offset_view.strides()[0], 3);
	EXPECT_EQ(offset_view.strides()[1], 1);

	// For CPU (direct address device), byte_offset should be added to data pointer
	// and byte_offset field should be 0
	// element_offset=1 for float32 = 4 bytes
	size_t expected_byte_offset =
	GetDataSize(static_cast<size_t>(element_offset), DLDataType({kDLFloat, 32, 1}));
	EXPECT_EQ(expected_byte_offset, 4); // 1 element * 32 bits / 8 = 4 bytes

	// The data pointer should be advanced by 4 bytes (1 float element)
	void* expected_offset_ptr = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset;
	EXPECT_EQ(offset_view.data_ptr(), expected_offset_ptr);
	EXPECT_EQ(offset_view.byte_offset(), 0); // Should be 0 for direct address devices

	// Verify data access through the offset view
	float* offset_data = reinterpret_cast<float*>(offset_view.data_ptr());
	EXPECT_EQ(offset_data[0 * 3 + 0 * 1], 1.0f); // Points to data[1]
	EXPECT_EQ(offset_data[1 * 3 + 0 * 1], 4.0f); // Points to data[4]

	// Test with larger element_offset
	int64_t element_offset2 = 2;
	Shape offset_shape2 = {1, 2};
	Shape offset_strides2 = {3, 1};
	TensorView offset_view2 = tensor_view.as_strided(offset_shape2, offset_strides2, element_offset2);
	size_t expected_byte_offset2 =
	GetDataSize(static_cast<size_t>(element_offset2), DLDataType({kDLFloat, 32, 1}));
	EXPECT_EQ(expected_byte_offset2, 8); // 2 elements * 32 bits / 8 = 8 bytes
	void* expected_offset_ptr2 = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset2;
	EXPECT_EQ(offset_view2.data_ptr(), expected_offset_ptr2);
	EXPECT_EQ(offset_view2.byte_offset(), 0);

	float* offset_data2 = reinterpret_cast<float*>(offset_view2.data_ptr());
	EXPECT_EQ(offset_data2[0 * 3 + 0 * 1], 2.0f); // Points to data[2]
	}

	TEST(Tensor, AsStrided) {
	// Create a base tensor with shape [2, 3] = 6 elements
	Tensor tensor = Empty({2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));

	// Fill with sequential values: [0, 1, 2, 3, 4, 5]
	float* data = reinterpret_cast<float*>(tensor.data_ptr());
	size_t element_capacity = GetDataSize(tensor) / sizeof(float);
	ASSERT_EQ(element_capacity, static_cast<size_t>(tensor.numel()));
	for (size_t i = 0; i < element_capacity; ++i) {
	data[i] = static_cast<float>(i);
	}

	void* original_data_ptr = tensor.data_ptr();
	EXPECT_EQ(tensor.byte_offset(), 0);

	// Create a strided view with shape [3, 2] and custom strides
	Shape new_shape = {3, 2};
	Shape new_strides = {1, 3};
	Tensor strided_view = tensor.as_strided(new_shape, new_strides);

	// Verify the view has correct shape and strides
	EXPECT_EQ(strided_view.shape().size(), 2);
	EXPECT_EQ(strided_view.shape()[0], 3);
	EXPECT_EQ(strided_view.shape()[1], 2);
	EXPECT_EQ(strided_view.strides().size(), 2);
	EXPECT_EQ(strided_view.strides()[0], 1);
	EXPECT_EQ(strided_view.strides()[1], 3);

	// Verify the view shares the same underlying data pointer (no offset)
	EXPECT_EQ(strided_view.data_ptr(), original_data_ptr);
	EXPECT_EQ(strided_view.byte_offset(), 0);
	EXPECT_EQ(strided_view.dtype(), tensor.dtype());

	// Test with element_offset - for float32, 1 element = 4 bytes
	Shape offset_shape = {2, 2};
	Shape offset_strides = {3, 1};
	int64_t element_offset = 1;
	Tensor offset_view = tensor.as_strided(offset_shape, offset_strides, element_offset);

	EXPECT_EQ(offset_view.shape().size(), 2);
	EXPECT_EQ(offset_view.shape()[0], 2);
	EXPECT_EQ(offset_view.shape()[1], 2);
	EXPECT_EQ(offset_view.strides().size(), 2);
	EXPECT_EQ(offset_view.strides()[0], 3);
	EXPECT_EQ(offset_view.strides()[1], 1);

	// For CPU (direct address device), byte_offset should be added to data pointer
	// and byte_offset field should be 0
	// element_offset=1 for float32 = 4 bytes
	size_t expected_byte_offset =
	GetDataSize(static_cast<size_t>(element_offset), DLDataType({kDLFloat, 32, 1}));
	EXPECT_EQ(expected_byte_offset, 4); // 1 element * 32 bits / 8 = 4 bytes

	// The data pointer should be advanced by 4 bytes (1 float element)
	void* expected_offset_ptr = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset;
	EXPECT_EQ(offset_view.data_ptr(), expected_offset_ptr);
	EXPECT_EQ(offset_view.byte_offset(), 0); // Should be 0 for direct address devices

	// Verify data access through the offset view
	float* offset_data = reinterpret_cast<float*>(offset_view.data_ptr());
	EXPECT_EQ(offset_data[0 * 3 + 0 * 1], 1.0f); // Points to data[1]
	EXPECT_EQ(offset_data[1 * 3 + 0 * 1], 4.0f); // Points to data[4]

	// Test with larger element_offset
	int64_t element_offset2 = 2;
	Tensor offset_view2 = tensor.as_strided({1, 2}, {3, 1}, element_offset2);
	size_t expected_byte_offset2 =
	GetDataSize(static_cast<size_t>(element_offset2), DLDataType({kDLFloat, 32, 1}));
	EXPECT_EQ(expected_byte_offset2, 8); // 2 elements * 32 bits / 8 = 8 bytes
	void* expected_offset_ptr2 = reinterpret_cast<char*>(original_data_ptr) + expected_byte_offset2;
	EXPECT_EQ(offset_view2.data_ptr(), expected_offset_ptr2);
	EXPECT_EQ(offset_view2.byte_offset(), 0);

	float* offset_data2 = reinterpret_cast<float*>(offset_view2.data_ptr());
	EXPECT_EQ(offset_data2[0 * 3 + 0 * 1], 2.0f); // Points to data[2]
	}

	TEST(Tensor, SizeStrideOutOfBounds) {
	Tensor tensor = Empty({2, 3, 4}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
	EXPECT_THROW({ tensor.size(3); }, tvm::ffi::Error);
	EXPECT_THROW({ tensor.size(-4); }, tvm::ffi::Error);
	EXPECT_THROW({ tensor.stride(3); }, tvm::ffi::Error);
	EXPECT_THROW({ tensor.stride(-4); }, tvm::ffi::Error);

	TensorView tensor_view = tensor;
	EXPECT_THROW({ tensor_view.size(3); }, tvm::ffi::Error);
	EXPECT_THROW({ tensor_view.size(-4); }, tvm::ffi::Error);
	EXPECT_THROW({ tensor_view.stride(3); }, tvm::ffi::Error);
	EXPECT_THROW({ tensor_view.stride(-4); }, tvm::ffi::Error);
	}

	} // namespace