add support in python api

Author: ZHUI

python/paddle/incubate/tensor/math.py

@@ -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.
@@ -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")
 
@@ -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.
@@ -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")
 
@@ -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.
@@ -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")
 
@@ -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.
@@ -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")