Incubator-mxnet: Can mx.nd.where(condition, x, y) supports if both x and y are None?

[huangzehao]

Hi, can mx.nd.where(condition, x, y) supports if both x and y are None, like tf.where and np.where?
If both x and y are None, then this operation returns the coordinates of true elements of condition.
@piiswrong @zhreshold
Thanks!

[zhreshold]

I think it's a good idea. We used to be stuck with the infer_shape of nd.where.
But I guess we'd better support it since so many ppl are asking for the functionality?

How about adding sentinel values to output so we can still infer the shape at this moment? @piiswrong

[zhreshold]

cc @szha @reminisce

[RogerChern]

Now we have mx.sym.scatter and mx.sym.gather in v1.0.0 which demand a fully implemented mx.sym.where. Please take this into consideration.

[ZiyueHuang]

This requires dynamic memory allocation and infer shape at runtime when we get the data. So some operators like tf.where, tf.unique, etc. cannot be implemented in MXNet currently.

Should we add support for dynamic shape? @piiswrong @reminisce

Following is the way in tensorflow, (please correct me if I'm wrong since I'm not familiar with tensorflow)

https://www.tensorflow.org/programmers_guide/faq

How can I determine the shape of a tensor in Python?
In TensorFlow, a tensor has both a static (inferred) shape and a dynamic (true) shape. The static shape can be read using the tf.Tensor.get_shape method: this shape is inferred from the operations that were used to create the tensor, and may be partially complete. If the static shape is not fully defined, the dynamic shape of a Tensor t can be determined by evaluating tf.shape(t).

[reminisce]

More and more people are asking for this feature. I think we should start considering adding this to MXNet. One way I can think of doing this is introducing a new ndarray storage type, skip shape inference for that type, and allocate the memory for it similar to sparse tensors in FCompute functions. The problem of this method is that it only applies to imperative mode and has to wait until the execution of the operator finishes to begin the execution of the next operator.
@eric-haibin-lin

[eric-haibin-lin]

@reminisce why a new storage type is required? What's the implication of that on existing operators and what operators will be supported for this new storage type? How will it interact with other operators?
Waiting until the completion of the execution is necessary, since we have to know the size of the output ndarray.

[eric-haibin-lin]

What we can do is:

• Add imperative support for nd.where. symbolic execution is not supported. To do so, we just need to rename the backend operator and add wrappers in the ndarray/symbol frontend, where ndarray.where supports x=None and y=None. It's hard to support symbolic execution since we have to block and wait on the completion of such operation, which makes shape inference of the rest of the operators in the graph very hard.
• Internally, the _where(cond) operator will register FComputeEx, and output csr output. The size of the output is [num_dims, N], where N is the number of total elements in cond, num_dims is the number of dimension of cond. At runtime, the operator will generate data based on n, the number of true elements in cond. To know how many true elements are there, we just need to access output.data. The size of output.data is (num_dim x n).
• Interface-wise, we should not return the csr output to user directly. Instead, the wrapper first call _where(cond) to get the csr output, and access output.data to get the indices. In order to make it easier for users to access these elements using advanced indexing, we can take slice the data to generate a list of num_dim 1-D ndarrays, where each one is of dimension n.

* currently advanced indexing only takes tuple of ndarrays as input indices. @reminisce it will be good to confirm if we can turn on the support of ndarray list easily (if there's no ambiguity).
* this is more like an ad-hot solution for where. If there are more general use cases like this, we should bring them up and discuss a general solution (like the dynamic shape in tf). If we want to go for a temporary solution that works and add dynamic shape later, we should make sure the interface is consistent from the user's perspective.

@szha @piiswrong @ZiyueHuang

[reminisce]

We can enable advanced indexing to support index as a list of numpy.ndarrays or mx.ndarray.NDArrays. They should have the same behavior as a tuple of ndarrays according to Numpy's behavior.

[Godricly]

Is there any update for this feature? Casting nd.array to numpy costs more time than np.where. :disappointed:

[eric-haibin-lin]

We can now follow the boolean mask example where dynamic output shape is supported:
src/operator/contrib/boolean_mask.cc