wehub-resource-sync
621 lines
22 KiB
C++
621 lines
22 KiB
C++
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed 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 <ATen/Functions.h>
|
|
#include <ATen/core/TensorBody.h>
|
|
#include <ATen/cuda/EmptyTensor.h>
|
|
#include <ATen/native/cuda/Resize.h>
|
|
#include <ATen/ops/tensor.h>
|
|
#include <c10/core/Layout.h>
|
|
#include <c10/core/ScalarType.h>
|
|
#include <c10/core/SymInt.h>
|
|
#include <c10/core/TensorOptions.h>
|
|
#include <vector>
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
#include <c10/cuda/CUDAFunctions.h>
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#endif
|
|
#include "ATen/ATen.h"
|
|
#include "gtest/gtest.h"
|
|
#include "paddle/phi/common/float16.h"
|
|
#include "torch/all.h"
|
|
|
|
// ============== at::zeros sparse tests ==============
|
|
|
|
TEST(SparseZerosTest, SparseCOO) {
|
|
// Dense tensor should return false
|
|
at::TensorBase tensor = at::zeros({2, 3}, at::kFloat);
|
|
ASSERT_FALSE(tensor.is_sparse());
|
|
ASSERT_EQ(tensor.layout(), c10::kStrided);
|
|
|
|
// Create sparse COO tensor
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparse);
|
|
at::TensorBase sparse_tensor = at::zeros({2, 3}, options);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
}
|
|
|
|
TEST(SparseZerosTest, SparseCsr) {
|
|
// Create sparse CSR tensor
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparseCsr);
|
|
at::TensorBase sparse_csr_tensor = at::zeros({2, 3}, options);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseZerosTest, WithOptionalParams) {
|
|
// Test zeros with optional parameters
|
|
at::Tensor sparse_tensor = at::zeros(
|
|
{2, 3}, at::kFloat, at::kSparse, at::kCPU, /*pin_memory=*/false);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
|
|
at::Tensor sparse_csr_tensor = at::zeros(
|
|
{2, 3}, at::kFloat, at::kSparseCsr, at::kCPU, /*pin_memory=*/false);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
// ============== at::empty sparse tests ==============
|
|
|
|
TEST(SparseEmptyTest, SparseCOO) {
|
|
// Dense tensor should return false
|
|
at::TensorBase tensor = at::empty({2, 3}, at::kFloat);
|
|
ASSERT_FALSE(tensor.is_sparse());
|
|
ASSERT_EQ(tensor.layout(), c10::kStrided);
|
|
|
|
// Create sparse COO tensor
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparse);
|
|
at::TensorBase sparse_tensor = at::empty({2, 3}, options);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
}
|
|
|
|
TEST(SparseEmptyTest, SparseCsr) {
|
|
// Create sparse CSR tensor
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparseCsr);
|
|
at::TensorBase sparse_csr_tensor = at::empty({2, 3}, options);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseEmptyTest, WithOptionalParams) {
|
|
// Test empty with optional parameters
|
|
at::Tensor sparse_tensor = at::empty({2, 3},
|
|
at::kFloat,
|
|
at::kSparse,
|
|
at::kCPU,
|
|
/*pin_memory=*/false,
|
|
/*memory_format=*/std::nullopt);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
|
|
at::Tensor sparse_csr_tensor = at::empty({2, 3},
|
|
at::kFloat,
|
|
at::kSparseCsr,
|
|
at::kCPU,
|
|
/*pin_memory=*/false,
|
|
/*memory_format=*/std::nullopt);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
// ============== at::empty_like sparse tests ==============
|
|
|
|
TEST(SparseEmptyLikeTest, SparseCOO) {
|
|
at::Tensor base_tensor = at::ones({2, 3}, at::kFloat);
|
|
|
|
// Dense empty_like should return false
|
|
at::TensorBase tensor = at::empty_like(base_tensor);
|
|
ASSERT_FALSE(tensor.is_sparse());
|
|
ASSERT_EQ(tensor.layout(), c10::kStrided);
|
|
|
|
// Create sparse COO tensor using empty_like
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparse);
|
|
at::TensorBase sparse_tensor = at::empty_like(base_tensor, options);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
}
|
|
|
|
TEST(SparseEmptyLikeTest, SparseCsr) {
|
|
at::Tensor base_tensor = at::ones({2, 3}, at::kFloat);
|
|
|
|
// Create sparse CSR tensor using empty_like
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparseCsr);
|
|
at::TensorBase sparse_csr_tensor = at::empty_like(base_tensor, options);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseEmptyLikeTest, WithOptionalParams) {
|
|
at::Tensor base_tensor = at::ones({2, 3}, at::kFloat);
|
|
|
|
// Test empty_like with optional parameters
|
|
at::Tensor sparse_tensor = at::empty_like(base_tensor,
|
|
at::kFloat,
|
|
at::kSparse,
|
|
at::kCPU,
|
|
/*pin_memory=*/false,
|
|
/*memory_format=*/std::nullopt);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
|
|
at::Tensor sparse_csr_tensor = at::empty_like(base_tensor,
|
|
at::kFloat,
|
|
at::kSparseCsr,
|
|
at::kCPU,
|
|
/*pin_memory=*/false,
|
|
/*memory_format=*/std::nullopt);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
// ============== at::zeros_like sparse tests ==============
|
|
|
|
TEST(SparseZerosLikeTest, SparseCOO) {
|
|
at::Tensor base_tensor = at::ones({2, 3}, at::kFloat);
|
|
|
|
// Dense zeros_like should return false
|
|
at::TensorBase tensor = at::zeros_like(base_tensor);
|
|
ASSERT_FALSE(tensor.is_sparse());
|
|
ASSERT_EQ(tensor.layout(), c10::kStrided);
|
|
|
|
// Create sparse COO tensor using zeros_like
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparse);
|
|
at::TensorBase sparse_tensor = at::zeros_like(base_tensor, options);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
}
|
|
|
|
TEST(SparseZerosLikeTest, SparseCsr) {
|
|
at::Tensor base_tensor = at::ones({2, 3}, at::kFloat);
|
|
|
|
// Create sparse CSR tensor using zeros_like
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparseCsr);
|
|
at::TensorBase sparse_csr_tensor = at::zeros_like(base_tensor, options);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseZerosLikeTest, WithOptionalParams) {
|
|
at::Tensor base_tensor = at::ones({2, 3}, at::kFloat);
|
|
|
|
// Test zeros_like with optional parameters
|
|
at::Tensor sparse_tensor = at::zeros_like(base_tensor,
|
|
at::kFloat,
|
|
at::kSparse,
|
|
at::kCPU,
|
|
/*pin_memory=*/false,
|
|
/*memory_format=*/std::nullopt);
|
|
ASSERT_TRUE(sparse_tensor.is_sparse());
|
|
ASSERT_EQ(sparse_tensor.layout(), c10::kSparse);
|
|
|
|
at::Tensor sparse_csr_tensor = at::zeros_like(base_tensor,
|
|
at::kFloat,
|
|
at::kSparseCsr,
|
|
at::kCPU,
|
|
/*pin_memory=*/false,
|
|
/*memory_format=*/std::nullopt);
|
|
ASSERT_TRUE(sparse_csr_tensor.is_sparse_csr());
|
|
ASSERT_EQ(sparse_csr_tensor.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
// ============== at::sparse_coo_tensor tests ==============
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensorBasic) {
|
|
// Create indices: 2D tensor of shape [sparse_dim, nnz]
|
|
// For a 3x4 sparse tensor with 2 non-zero elements at (0,1) and (2,3)
|
|
at::Tensor indices = at::empty({2, 2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
indices_ptr[0] = 0; // row index of first non-zero
|
|
indices_ptr[1] = 2; // row index of second non-zero
|
|
indices_ptr[2] = 1; // col index of first non-zero
|
|
indices_ptr[3] = 3; // col index of second non-zero
|
|
|
|
// Create values: tensor of shape [nnz]
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 1.0f;
|
|
values_ptr[1] = 2.0f;
|
|
|
|
// Create sparse COO tensor
|
|
at::Tensor sparse = at::sparse_coo_tensor(indices, values, {3, 4});
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensorWithOptions) {
|
|
at::Tensor indices = at::empty({2, 2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
indices_ptr[0] = 0;
|
|
indices_ptr[1] = 1;
|
|
indices_ptr[2] = 0;
|
|
indices_ptr[3] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 3.0f;
|
|
values_ptr[1] = 4.0f;
|
|
|
|
// Create with optional parameters
|
|
at::Tensor sparse = at::sparse_coo_tensor(indices,
|
|
values,
|
|
{2, 2},
|
|
at::kFloat,
|
|
at::kSparse,
|
|
at::kCPU,
|
|
/*pin_memory=*/false,
|
|
std::nullopt);
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensorWithCoalescedOptionTrue) {
|
|
at::Tensor indices = at::empty({2, 2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
indices_ptr[0] = 0;
|
|
indices_ptr[1] = 1;
|
|
indices_ptr[2] = 0;
|
|
indices_ptr[3] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
values.data_ptr<float>()[0] = 3.0f;
|
|
values.data_ptr<float>()[1] = 4.0f;
|
|
|
|
at::Tensor sparse =
|
|
at::sparse_coo_tensor(indices, values, {2, 2}, at::TensorOptions(), true);
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
ASSERT_TRUE(sparse.is_coalesced());
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensorWithCoalescedOptionFalse) {
|
|
at::Tensor indices = at::empty({2, 2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
indices_ptr[0] = 0;
|
|
indices_ptr[1] = 1;
|
|
indices_ptr[2] = 0;
|
|
indices_ptr[3] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
values.data_ptr<float>()[0] = 3.0f;
|
|
values.data_ptr<float>()[1] = 4.0f;
|
|
|
|
at::Tensor sparse = at::sparse_coo_tensor(
|
|
indices, values, {2, 2}, at::TensorOptions(), false);
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
ASSERT_FALSE(sparse.is_coalesced());
|
|
}
|
|
|
|
// ============== at::sparse_csr_tensor tests ==============
|
|
|
|
TEST(SparseConstructorTest, SparseCsrTensorBasic) {
|
|
// Create a 3x4 sparse CSR tensor with 4 non-zero elements:
|
|
// Row 0: values at columns 0, 2
|
|
// Row 1: value at column 1
|
|
// Row 2: value at column 3
|
|
|
|
// crow_indices: [0, 2, 3, 4] - compressed row pointers
|
|
at::Tensor crow_indices =
|
|
at::empty({4}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* crow_ptr = crow_indices.data_ptr<int64_t>();
|
|
crow_ptr[0] = 0;
|
|
crow_ptr[1] = 2;
|
|
crow_ptr[2] = 3;
|
|
crow_ptr[3] = 4;
|
|
|
|
// col_indices: [0, 2, 1, 3] - column indices
|
|
at::Tensor col_indices =
|
|
at::empty({4}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* col_ptr = col_indices.data_ptr<int64_t>();
|
|
col_ptr[0] = 0;
|
|
col_ptr[1] = 2;
|
|
col_ptr[2] = 1;
|
|
col_ptr[3] = 3;
|
|
|
|
// values: [1.0, 2.0, 3.0, 4.0]
|
|
at::Tensor values = at::empty({4}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 1.0f;
|
|
values_ptr[1] = 2.0f;
|
|
values_ptr[2] = 3.0f;
|
|
values_ptr[3] = 4.0f;
|
|
|
|
// Create sparse CSR tensor
|
|
at::Tensor sparse = at::sparse_csr_tensor(
|
|
crow_indices, col_indices, values, {3, 4}, at::TensorOptions());
|
|
|
|
ASSERT_TRUE(sparse.is_sparse_csr());
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCsrTensorWithOptions) {
|
|
// Create a simple 2x2 sparse CSR tensor
|
|
at::Tensor crow_indices =
|
|
at::empty({3}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* crow_ptr = crow_indices.data_ptr<int64_t>();
|
|
crow_ptr[0] = 0;
|
|
crow_ptr[1] = 1;
|
|
crow_ptr[2] = 2;
|
|
|
|
at::Tensor col_indices =
|
|
at::empty({2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* col_ptr = col_indices.data_ptr<int64_t>();
|
|
col_ptr[0] = 0;
|
|
col_ptr[1] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 5.0f;
|
|
values_ptr[1] = 6.0f;
|
|
|
|
// Create with optional parameters
|
|
at::Tensor sparse = at::sparse_csr_tensor(crow_indices,
|
|
col_indices,
|
|
values,
|
|
{2, 2},
|
|
at::kFloat,
|
|
at::kSparseCsr,
|
|
at::kCPU,
|
|
/*pin_memory=*/false);
|
|
|
|
ASSERT_TRUE(sparse.is_sparse_csr());
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCsrTensorMismatchedOptionsDtypeIgnored) {
|
|
// PyTorch ignores dtype mismatch in sparse_csr_tensor;
|
|
// the resulting tensor uses values' original dtype.
|
|
at::Tensor crow_indices =
|
|
at::empty({3}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* crow_ptr = crow_indices.data_ptr<int64_t>();
|
|
crow_ptr[0] = 0;
|
|
crow_ptr[1] = 1;
|
|
crow_ptr[2] = 2;
|
|
|
|
at::Tensor col_indices =
|
|
at::empty({2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* col_ptr = col_indices.data_ptr<int64_t>();
|
|
col_ptr[0] = 0;
|
|
col_ptr[1] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 5.0f;
|
|
values_ptr[1] = 6.0f;
|
|
|
|
std::vector<int64_t> size = {2, 2};
|
|
auto options = c10::TensorOptions().dtype(at::kDouble);
|
|
|
|
at::Tensor sparse =
|
|
at::sparse_csr_tensor(crow_indices, col_indices, values, size, options);
|
|
|
|
// Result should use values' dtype (float), not options' dtype (double).
|
|
ASSERT_EQ(sparse.dtype(), at::kFloat);
|
|
}
|
|
|
|
// ============== Additional sparse_coo_tensor tests ==============
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensorInferSize) {
|
|
// Test sparse_coo_tensor with inferred size (no explicit size parameter)
|
|
at::Tensor indices = at::empty({2, 3}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
// 3 non-zero elements at (0,0), (1,1), (2,2)
|
|
indices_ptr[0] = 0;
|
|
indices_ptr[1] = 1;
|
|
indices_ptr[2] = 2;
|
|
indices_ptr[3] = 0;
|
|
indices_ptr[4] = 1;
|
|
indices_ptr[5] = 2;
|
|
|
|
at::Tensor values = at::empty({3}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 1.0f;
|
|
values_ptr[1] = 2.0f;
|
|
values_ptr[2] = 3.0f;
|
|
|
|
// Create sparse COO tensor with inferred size
|
|
at::Tensor sparse = at::sparse_coo_tensor(indices, values);
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
ASSERT_EQ(sparse.dim(), 2);
|
|
ASSERT_EQ(sparse.size(0), 3);
|
|
ASSERT_EQ(sparse.size(1), 3);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensorInferSizeWithCoalescedOption) {
|
|
at::Tensor indices = at::empty({2, 2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
indices_ptr[0] = 0;
|
|
indices_ptr[1] = 2;
|
|
indices_ptr[2] = 1;
|
|
indices_ptr[3] = 3;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
values.data_ptr<float>()[0] = 1.0f;
|
|
values.data_ptr<float>()[1] = 2.0f;
|
|
|
|
at::Tensor sparse =
|
|
at::sparse_coo_tensor(indices, values, at::TensorOptions(), true);
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
ASSERT_EQ(sparse.size(0), 3);
|
|
ASSERT_EQ(sparse.size(1), 4);
|
|
ASSERT_TRUE(sparse.is_coalesced());
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensorDouble) {
|
|
// Test sparse_coo_tensor with double dtype
|
|
at::Tensor indices = at::empty({2, 2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
indices_ptr[0] = 0;
|
|
indices_ptr[1] = 1;
|
|
indices_ptr[2] = 0;
|
|
indices_ptr[3] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kDouble));
|
|
double* values_ptr = values.data_ptr<double>();
|
|
values_ptr[0] = 1.5;
|
|
values_ptr[1] = 2.5;
|
|
|
|
at::Tensor sparse = at::sparse_coo_tensor(indices, values, {2, 2});
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCooTensor3D) {
|
|
// Test 3D sparse COO tensor
|
|
at::Tensor indices = at::empty({3, 2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* indices_ptr = indices.data_ptr<int64_t>();
|
|
// 2 non-zero elements at (0,1,2) and (1,0,1)
|
|
indices_ptr[0] = 0;
|
|
indices_ptr[1] = 1;
|
|
indices_ptr[2] = 1;
|
|
indices_ptr[3] = 0;
|
|
indices_ptr[4] = 2;
|
|
indices_ptr[5] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 5.0f;
|
|
values_ptr[1] = 6.0f;
|
|
|
|
at::Tensor sparse = at::sparse_coo_tensor(indices, values, {2, 2, 3});
|
|
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparse);
|
|
}
|
|
|
|
// ============== Additional sparse_csr_tensor tests ==============
|
|
|
|
TEST(SparseConstructorTest, SparseCsrTensorDouble) {
|
|
// Test sparse_csr_tensor with double dtype
|
|
at::Tensor crow_indices =
|
|
at::empty({3}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* crow_ptr = crow_indices.data_ptr<int64_t>();
|
|
crow_ptr[0] = 0;
|
|
crow_ptr[1] = 1;
|
|
crow_ptr[2] = 2;
|
|
|
|
at::Tensor col_indices =
|
|
at::empty({2}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* col_ptr = col_indices.data_ptr<int64_t>();
|
|
col_ptr[0] = 0;
|
|
col_ptr[1] = 1;
|
|
|
|
at::Tensor values = at::empty({2}, c10::TensorOptions().dtype(at::kDouble));
|
|
double* values_ptr = values.data_ptr<double>();
|
|
values_ptr[0] = 1.5;
|
|
values_ptr[1] = 2.5;
|
|
|
|
at::Tensor sparse = at::sparse_csr_tensor(
|
|
crow_indices, col_indices, values, {2, 2}, at::TensorOptions());
|
|
|
|
ASSERT_TRUE(sparse.is_sparse_csr());
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCsrTensorLarger) {
|
|
// Test larger sparse CSR tensor (4x5 with 6 non-zero elements)
|
|
// Row 0: values at columns 1, 3
|
|
// Row 1: value at column 2
|
|
// Row 2: values at columns 0, 4
|
|
// Row 3: value at column 2
|
|
|
|
at::Tensor crow_indices =
|
|
at::empty({5}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* crow_ptr = crow_indices.data_ptr<int64_t>();
|
|
crow_ptr[0] = 0;
|
|
crow_ptr[1] = 2;
|
|
crow_ptr[2] = 3;
|
|
crow_ptr[3] = 5;
|
|
crow_ptr[4] = 6;
|
|
|
|
at::Tensor col_indices =
|
|
at::empty({6}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* col_ptr = col_indices.data_ptr<int64_t>();
|
|
col_ptr[0] = 1;
|
|
col_ptr[1] = 3;
|
|
col_ptr[2] = 2;
|
|
col_ptr[3] = 0;
|
|
col_ptr[4] = 4;
|
|
col_ptr[5] = 2;
|
|
|
|
at::Tensor values = at::empty({6}, c10::TensorOptions().dtype(at::kFloat));
|
|
float* values_ptr = values.data_ptr<float>();
|
|
values_ptr[0] = 1.0f;
|
|
values_ptr[1] = 2.0f;
|
|
values_ptr[2] = 3.0f;
|
|
values_ptr[3] = 4.0f;
|
|
values_ptr[4] = 5.0f;
|
|
values_ptr[5] = 6.0f;
|
|
|
|
at::Tensor sparse = at::sparse_csr_tensor(
|
|
crow_indices, col_indices, values, {4, 5}, at::TensorOptions());
|
|
|
|
ASSERT_TRUE(sparse.is_sparse_csr());
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
TEST(SparseConstructorTest, SparseCsrTensorEmpty) {
|
|
// Test sparse CSR tensor with no non-zero elements
|
|
at::Tensor crow_indices =
|
|
at::empty({4}, c10::TensorOptions().dtype(at::kLong));
|
|
int64_t* crow_ptr = crow_indices.data_ptr<int64_t>();
|
|
crow_ptr[0] = 0;
|
|
crow_ptr[1] = 0;
|
|
crow_ptr[2] = 0;
|
|
crow_ptr[3] = 0;
|
|
|
|
at::Tensor col_indices =
|
|
at::empty({0}, c10::TensorOptions().dtype(at::kLong));
|
|
|
|
at::Tensor values = at::empty({0}, c10::TensorOptions().dtype(at::kFloat));
|
|
|
|
at::Tensor sparse = at::sparse_csr_tensor(
|
|
crow_indices, col_indices, values, {3, 3}, at::TensorOptions());
|
|
|
|
ASSERT_TRUE(sparse.is_sparse_csr());
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparseCsr);
|
|
}
|
|
|
|
// ============== Sparse tensor interoperability tests ==============
|
|
|
|
TEST(SparseInteropTest, SparseCsrFromZeros) {
|
|
// Create sparse CSR tensor from zeros
|
|
auto options = c10::TensorOptions().dtype(at::kFloat).layout(at::kSparseCsr);
|
|
at::Tensor sparse = at::zeros({4, 4}, options);
|
|
|
|
ASSERT_TRUE(sparse.is_sparse_csr());
|
|
ASSERT_TRUE(sparse.is_sparse());
|
|
ASSERT_EQ(sparse.layout(), c10::kSparseCsr);
|
|
}
|