Skip to content

Commit 121745f

Browse files
NKNaNNKNaN
authored
【Hackathon 7th No.25】为 Paddle 新增 sparse_dim/dense_dim -part (#69132)
* add sparse dim dense dim * fix example code * fix test * resolve conflict * update pir
1 parent cd3d988 commit 121745f

File tree

7 files changed

+797
-3
lines changed

7 files changed

+797
-3
lines changed

‎paddle/fluid/pybind/eager_method.cc

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3316,6 +3316,121 @@ static PyObject* tensor_is_contiguous(TensorObject* self,
33163316
EAGER_CATCH_AND_THROW_RETURN_NULL
33173317
}
33183318

3319+
PyDoc_STRVAR(tensor_method_sparse_dim__doc__,
3320+
R"DOC(sparse_dim($self, /)
3321+
--
3322+
3323+
Returns the number of sparse dimensions of sparse Tensor.
3324+
3325+
Note:
3326+
**If self is not sparse Tensor, return 0.**
3327+
3328+
Returns:
3329+
int, sparse dim of self Tensor
3330+
3331+
Examples:
3332+
3333+
.. code-block:: python
3334+
3335+
>>> import paddle
3336+
3337+
>>> indices = [[0, 1, 2], [1, 2, 0]]
3338+
>>> values = [1.0, 2.0, 3.0]
3339+
>>> dense_shape = [3, 3]
3340+
>>> coo = paddle.sparse.sparse_coo_tensor(indices, values, dense_shape)
3341+
>>> coo.sparse_dim()
3342+
2
3343+
3344+
>>> crows = [0, 2, 3, 5]
3345+
>>> cols = [1, 3, 2, 0, 1]
3346+
>>> values = [1, 2, 3, 4, 5]
3347+
>>> dense_shape = [3, 4]
3348+
>>> csr = paddle.sparse.sparse_csr_tensor(crows, cols, values, dense_shape)
3349+
>>> csr.sparse_dim()
3350+
2
3351+
3352+
>>> dense = paddle.to_tensor([1, 2, 3])
3353+
>>> dense.sparse_dim()
3354+
0
3355+
3356+
)DOC"); // NOLINT
3357+
3358+
static PyObject* tensor_method_sparse_dim(TensorObject* self,
3359+
PyObject* args,
3360+
PyObject* kwargs) {
3361+
EAGER_TRY
3362+
if (self->tensor.is_sparse_coo_tensor()) {
3363+
auto sparse_coo_tensor =
3364+
std::dynamic_pointer_cast<phi::SparseCooTensor>(self->tensor.impl());
3365+
return ToPyObject(sparse_coo_tensor->sparse_dim());
3366+
} else if (self->tensor.is_sparse_csr_tensor()) {
3367+
auto sparse_csr_tensor =
3368+
std::dynamic_pointer_cast<phi::SparseCsrTensor>(self->tensor.impl());
3369+
return ToPyObject(sparse_csr_tensor->sparse_dim());
3370+
} else {
3371+
return ToPyObject(0);
3372+
}
3373+
EAGER_CATCH_AND_THROW_RETURN_NULL
3374+
}
3375+
3376+
PyDoc_STRVAR(tensor_method_dense_dim__doc__,
3377+
R"DOC(dense_dim($self, /)
3378+
--
3379+
3380+
Returns the number of dense dimensions of sparse Tensor.
3381+
3382+
Note:
3383+
**If self is not sparse Tensor, return len(self.shape).**
3384+
3385+
Returns:
3386+
int, dense dim of self Tensor
3387+
3388+
Examples:
3389+
3390+
.. code-block:: python
3391+
3392+
>>> import paddle
3393+
>>> import numpy as np
3394+
3395+
>>> indices = [[0, 1, 1], [2, 0, 2]]
3396+
>>> values = np.array([[3, 4], [5, 6], [7, 8]])
3397+
>>> dense_shape = [2, 3, 2]
3398+
>>> coo = paddle.sparse.sparse_coo_tensor(indices, values, dense_shape)
3399+
>>> coo.dense_dim()
3400+
1
3401+
3402+
>>> crows = [0, 2, 3, 5]
3403+
>>> cols = [1, 3, 2, 0, 1]
3404+
>>> values = [1, 2, 3, 4, 5]
3405+
>>> dense_shape = [3, 4]
3406+
>>> csr = paddle.sparse.sparse_csr_tensor(crows, cols, values, dense_shape)
3407+
>>> csr.dense_dim()
3408+
0
3409+
3410+
>>> dense = paddle.to_tensor([[1, 2, 3]])
3411+
>>> dense.dense_dim()
3412+
>>> 2
3413+
3414+
)DOC"); // NOLINT
3415+
3416+
static PyObject* tensor_method_dense_dim(TensorObject* self,
3417+
PyObject* args,
3418+
PyObject* kwargs) {
3419+
EAGER_TRY
3420+
if (self->tensor.is_sparse_coo_tensor()) {
3421+
auto sparse_coo_tensor =
3422+
std::dynamic_pointer_cast<phi::SparseCooTensor>(self->tensor.impl());
3423+
return ToPyObject(sparse_coo_tensor->dense_dim());
3424+
} else if (self->tensor.is_sparse_csr_tensor()) {
3425+
auto sparse_csr_tensor =
3426+
std::dynamic_pointer_cast<phi::SparseCsrTensor>(self->tensor.impl());
3427+
return ToPyObject(sparse_csr_tensor->dense_dim());
3428+
} else {
3429+
return ToPyObject(self->tensor.shape().size());
3430+
}
3431+
EAGER_CATCH_AND_THROW_RETURN_NULL
3432+
}
3433+
33193434
static PyObject* tensor_method__set_impl(TensorObject* self,
33203435
PyObject* args,
33213436
PyObject* kwargs) {
@@ -3581,6 +3696,14 @@ PyMethodDef variable_methods[] = { // NOLINT
35813696
(PyCFunction)(void (*)())tensor_method_is_coalesced,
35823697
METH_VARARGS | METH_KEYWORDS,
35833698
tensor_is_coalesced__doc__},
3699+
{"sparse_dim",
3700+
(PyCFunction)(void (*)())tensor_method_sparse_dim,
3701+
METH_VARARGS | METH_KEYWORDS,
3702+
tensor_method_sparse_dim__doc__},
3703+
{"dense_dim",
3704+
(PyCFunction)(void (*)())tensor_method_dense_dim,
3705+
METH_VARARGS | METH_KEYWORDS,
3706+
tensor_method_dense_dim__doc__},
35843707
/***the method of sparse tensor****/
35853708
{"element_size",
35863709
(PyCFunction)(void (*)())tensor_method_element_size,

‎paddle/fluid/pybind/pir.cc

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1461,9 +1461,42 @@ void BindValue(py::module *m) {
14611461
return py::cast<py::none>(Py_None);
14621462
}
14631463
})
1464-
.def("_clone", [](Value self) {
1465-
// Return a new value owned by python side
1466-
return self;
1464+
.def("_clone",
1465+
[](Value self) {
1466+
// Return a new value owned by python side
1467+
return self;
1468+
})
1469+
.def("sparse_dim",
1470+
[](Value self) -> int32_t {
1471+
auto op_result = self.dyn_cast<OpResult>();
1472+
pir::Operation *operation = op_result.owner();
1473+
if (self.type().isa<SparseCooTensorType>() &&
1474+
operation->name() == "pd_op.sparse_coo_tensor_sp") {
1475+
std::vector<Value> sources = operation->operands_source();
1476+
Value non_zero_indices = sources[1];
1477+
return phi::vectorize(GetValueDims(non_zero_indices))[0];
1478+
} else if (self.type().isa<SparseCsrTensorType>()) {
1479+
PADDLE_THROW(common::errors::InvalidType(
1480+
"SparseCsrTensor is unsupported in pir mode."));
1481+
} else {
1482+
return 0;
1483+
}
1484+
})
1485+
.def("dense_dim", [](Value self) -> int32_t {
1486+
auto op_result = self.dyn_cast<OpResult>();
1487+
pir::Operation *operation = op_result.owner();
1488+
if (self.type().isa<SparseCooTensorType>() &&
1489+
operation->name() == "pd_op.sparse_coo_tensor_sp") {
1490+
std::vector<Value> sources = operation->operands_source();
1491+
Value non_zero_indices = sources[1];
1492+
int32_t dims = phi::vectorize(GetValueDims(self)).size();
1493+
return dims - phi::vectorize(GetValueDims(non_zero_indices))[0];
1494+
} else if (self.type().isa<SparseCsrTensorType>()) {
1495+
PADDLE_THROW(common::errors::InvalidType(
1496+
"SparseCsrTensor is unsupported in pir mode."));
1497+
} else {
1498+
return phi::vectorize(GetValueDims(self)).size();
1499+
}
14671500
});
14681501
}
14691502

‎paddle/phi/core/sparse_csr_tensor.cc

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,4 +160,16 @@ void SparseCsrTensor::set_meta(const SparseTensorMeta& meta) {
160160
meta_.dtype = meta.dtype;
161161
meta_.layout = meta.layout;
162162
}
163+
164+
int32_t SparseCsrTensor::sparse_dim() const { return 2; }
165+
166+
int32_t SparseCsrTensor::dense_dim() const {
167+
int32_t nze_dim = this->non_zero_elements_.dims().size();
168+
int32_t batch_dim = this->non_zero_crows_.dims().size() - 1;
169+
// layout of SparseCsrTensor has not been implemented yet
170+
// int32_t block_dim = = (layout_ == kSparseBsr || layout_ == kSparseBsc ? 2
171+
// : 0);
172+
int32_t block_dim = 0;
173+
return nze_dim - batch_dim - block_dim - 1;
174+
}
163175
} // namespace phi

‎paddle/phi/core/sparse_csr_tensor.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,14 @@ class SparseCsrTensor : public TensorBase,
110110
/// \return The data type of the tensor.
111111
DataType dtype() const noexcept override { return meta_.dtype; }
112112

113+
/// \brief get the sparse dim
114+
/// \return The sparse dim of the tensor.
115+
int32_t sparse_dim() const;
116+
117+
/// \brief get the dense dim
118+
/// \return The dense dim of the tensor.
119+
int32_t dense_dim() const;
120+
113121
#ifndef PADDLE_WITH_CUSTOM_KERNEL
114122
void set_type(const DataType dtype);
115123
#endif

‎python/paddle/tensor/tensor.prototype.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,7 @@ class AbstractTensor:
225225
@data.setter
226226
def data(self, value: Tensor) -> None: ...
227227
def data_ptr(self) -> int: ...
228+
def dense_dim(self) -> int: ...
228229
def detach(self) -> Tensor: ...
229230
def detach_(self) -> Tensor: ...
230231
@property
@@ -286,6 +287,7 @@ class AbstractTensor:
286287
def shape(self) -> list[int]: ...
287288
@property
288289
def size(self) -> int: ...
290+
def sparse_dim(self) -> int: ...
289291
@property
290292
def stop_gradient(self) -> bool: ...
291293
@stop_gradient.setter

0 commit comments

Comments
 (0)