Skip to content

Commit

Permalink
add support in python api
Browse files Browse the repository at this point in the history
  • Loading branch information
@ZHUI
ZHUI committed Mar 16, 2022
1 parent c2771ec commit 880a6f2
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions python/paddle/incubate/tensor/math.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ def segment_sum(data, segment_ids, name=None):
where sum is over j such that `segment_ids[j] == i`.
Args:
data (Tensor): A tensor, available data type float32, float64.
data (Tensor): A tensor, available data type float32, float64, int32, int64.
segment_ids (Tensor): A 1-D tensor, which have the same size
with the first dimension of input data.
Available data type is int32, int64.
Expand All @@ -54,7 +54,8 @@ def segment_sum(data, segment_ids, name=None):
out, tmp = _C_ops.segment_pool(data, segment_ids, 'pooltype', "SUM")
return out

check_variable_and_dtype(data, "X", ("float32", "float64"), "segment_pool")
check_variable_and_dtype(data, "X", ("float32", "float64", "int32",
"int64"), "segment_pool")
check_variable_and_dtype(segment_ids, "SegmentIds", ("int32", "int64"),
"segment_pool")

Expand Down Expand Up @@ -82,7 +83,7 @@ def segment_mean(data, segment_ids, name=None):
of all index 'segment_ids[j] == i'.
Args:
data (tensor): a tensor, available data type float32, float64.
data (tensor): a tensor, available data type float32, float64, int32, int64.
segment_ids (tensor): a 1-d tensor, which have the same size
with the first dimension of input data.
available data type is int32, int64.
Expand All @@ -107,7 +108,8 @@ def segment_mean(data, segment_ids, name=None):
out, tmp = _C_ops.segment_pool(data, segment_ids, 'pooltype', "MEAN")
return out

check_variable_and_dtype(data, "X", ("float32", "float64"), "segment_pool")
check_variable_and_dtype(data, "X", ("float32", "float64", "int32",
"int64"), "segment_pool")
check_variable_and_dtype(segment_ids, "SegmentIds", ("int32", "int64"),
"segment_pool")

Expand All @@ -134,7 +136,7 @@ def segment_min(data, segment_ids, name=None):
where min is over j such that `segment_ids[j] == i`.
Args:
data (tensor): a tensor, available data type float32, float64.
data (tensor): a tensor, available data type float32, float64, int32, int64.
segment_ids (tensor): a 1-d tensor, which have the same size
with the first dimension of input data.
available data type is int32, int64.
Expand All @@ -159,7 +161,8 @@ def segment_min(data, segment_ids, name=None):
out, tmp = _C_ops.segment_pool(data, segment_ids, 'pooltype', "MIN")
return out

check_variable_and_dtype(data, "X", ("float32", "float64"), "segment_pool")
check_variable_and_dtype(data, "X", ("float32", "float64", "int32",
"int64"), "segment_pool")
check_variable_and_dtype(segment_ids, "SegmentIds", ("int32", "int64"),
"segment_pool")

Expand All @@ -186,7 +189,7 @@ def segment_max(data, segment_ids, name=None):
where max is over j such that `segment_ids[j] == i`.
Args:
data (tensor): a tensor, available data type float32, float64.
data (tensor): a tensor, available data type float32, float64, int32, int64.
segment_ids (tensor): a 1-d tensor, which have the same size
with the first dimension of input data.
available data type is int32, int64.
Expand All @@ -211,7 +214,8 @@ def segment_max(data, segment_ids, name=None):
out, tmp = _C_ops.segment_pool(data, segment_ids, 'pooltype', "MAX")
return out

check_variable_and_dtype(data, "X", ("float32", "float64"), "segment_pool")
check_variable_and_dtype(data, "X", ("float32", "float64", "int32",
"int64"), "segment_pool")
check_variable_and_dtype(segment_ids, "SegmentIds", ("int32", "int64"),
"segment_pool")

Expand Down

1 comment on commit 880a6f2

@paddle-bot-old
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅Congratulation! Your pull request passed all required CI. You could ask reviewer(s) to approve and merge. 🎉

Please sign in to comment.