def construct(self, a, b, indices, values):
x = RowTensor(indices, values, self.dense_shape)
x = self.op2(x)
while a > b:
x = self.op1(x)
b = b + 1
return x.indices, x.values, x.dense_shape
def construct(self, a, b, indices, values):
x = RowTensor(indices, values, self.dense_shape)
if a > b:
x = self.op1(x)
else:
x = self.op2(x)
return x.indices, x.values
def bprop(x, indices, axis, out, dout):
x_shp = shape_op(x)
if axis == 0:
indices_size = (size_op(indices), )
x_tail_shp = x_shp[1:]
values_shape = indices_size + x_tail_shp
values = reshape(dout, values_shape)
indices = reshape(indices, indices_size)
return RowTensor(indices, values,
x_shp), zeros_like(indices), zeros_like(axis)
if F.rank(dout) == 0:
dout = P.ExpandDims()(dout, -1)
if F.rank(indices) == 0:
indices = P.ExpandDims()(indices, -1)
out_shp = shape_op(dout)
ind_shp = shape_op(indices)
# Example: out_shape:(3,2,3) axis 1 -> (1,0,2)
perm_1 = _generate_shape_index(out_shp, ind_shp, axis)
values_transpose = transpose(dout, perm_1)
params_grad = unsorted_segment_sum(values_transpose, indices,
shape_op(x)[axis])
# Example: out_shape:(3,2,3) axis 2 -> (1,2,0)
perm_2 = _generate_inverse_index(x_shp, axis)
params_grad = transpose(params_grad, perm_2)
return params_grad, zeros_like(indices), zeros_like(axis)
def construct(self, x):
indices = x.indices
values = x.values + 2
dense_shape = x.dense_shape
return RowTensor(indices, values, dense_shape)
def construct(self, indices, values):
x = RowTensor(indices, values, self.dense_shape)
return x.values, x.indices, x.dense_shape
def construct(self, indices, values):
ret = (RowTensor(indices, values, self.dense_shape), )
return ret[0]