eric-haibin-lin commented on a change in pull request #10374: [MXNET-93] Sparse
support for Custom Op
URL: https://github.com/apache/incubator-mxnet/pull/10374#discussion_r180299395
##
File path: python/mxnet/operator.py
##
@@ -522,6 +526,69 @@ def infer_type(self, in_type):
return in_type, [in_type[0]]*len(self.list_outputs()), \
[in_type[0]]*len(self.list_auxiliary_states())
+def infer_storage_type(self, in_stype):
+"""infer_storage_type interface. Used to infer storage type of
+inputs and outputs in the forward pass.
+
+Parameters
+--
+in_stype : list of stypes, Valid stypes are default, row_sparse and
+csr
+
+Returns
+---
+in_stype : list
+list of argument stypes.
+out_stype : list
+list of output types calculated from in_stype,
+in the same order as declared in list_outputs.
+aux_type : Optional, list
+list of aux types calculated from in_stype,
+in the same order as declared in list_auxiliary_states.
+"""
+return in_stype, [in_stype[0]]*len(self.list_outputs()), \
+[in_stype[0]]*len(self.list_auxiliary_states())
+
+def infer_storage_type_backward(self, ograd_stype, in_stype, out_stype,
igrad_stype, aux_stype):
+"""infer_storage_type_backward interface. Used to infer storage
+type of inputs and outputs in the backward pass.
+
+Will raise an error if undefined storage type is returned.
+Returned lists have to be the same size as the input lists to
infer_storage_type_backward,
+otherwise an exception will be thrown. When this interface is not
implemented,
+all stypes will fallback to default.
Review comment:
I will be careful about using the word "fallback" here since it has specific
meaning for sparse ops.
I'm a little bit confused about the default behavior of forward stype
inference vs. backward stype inference.
In forward you replicated the stype of in_stype to outputs. In backward you
replicated the "default" stype to all outputs.
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
