未验证 提交 4681f13b 编写于 作者: Z zhangkaihuo 提交者: GitHub

deliver indices_dict (#46919)

上级 6f7af148
......@@ -169,6 +169,7 @@ void CoalesceCooGPUKernel(const GPUContext& dev_ctx,
indexs_ptr, const_dims, out_nnz, sparse_dim, out_indices.data<IntT>());
out->SetMember(out_indices, out_values, x.dims(), true);
out->SetIndicesDict(x.GetIndicesDict());
}
template <typename T, typename Context>
......
......@@ -37,6 +37,7 @@ namespace sparse {
EmptyLikeCooKernel<T, Context>(dev_ctx, x, out); \
phi::prefix##Kernel<T, Context>( \
dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements()); \
out->SetIndicesDict(x.GetIndicesDict()); \
} \
\
template <typename T, typename Context> \
......@@ -105,6 +106,7 @@ void ScaleCooKernel(const Context& dev_ctx,
bias,
bias_after_scale,
out->mutable_non_zero_elements());
out->SetIndicesDict(x.GetIndicesDict());
}
template <typename T, typename Context>
......@@ -155,6 +157,7 @@ void CastCooKernel(const Context& dev_ctx,
meta.set_dtype(value_dtype);
phi::CastKernel<T, Context>(dev_ctx, x_values, value_dtype, out_values);
}
out->SetIndicesDict(x.GetIndicesDict());
}
template <typename T, typename Context>
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册