[numpy] Support zero-dim and zero-size tensors in MXNet

[reminisce]

Description

This PR provides the infrastructure of supporting zero-dim and zero-size tensors as the first outcome of the initiative of introducing NumPy compatible coding experience into MXNet (see this RFC). Great thanks to many folks who have contributed design, implementation and code review to this feature. The credits should go to all of them. Sorry if I missed anyone below. It would be impossible to make so many things correct scattered all over the codebase with just a couple of hands.

• C++ backend and Python: @junrushao1994 @eric-haibin-lin @szha @zheng-da @wkcn @mli
• MKLDNN: @pengzhao-intel @TaoLv @lihaofd @ZhennanQin
• Scala: @yzhliu
• R: @hetong007
• Perl: @sergeykolychev

FAQ

What are zero-dim and zero-size tensors?

• Zero-dim tensors are scalars with shape equal to ().
• A zero-size tensor is the one whose shape has at least one dimension size as 0. For example, np.ones((1, 0, 3)) would generate a zero-size tensor with shape equal to (1, 0, 3).

Why are they important?

Mathematically speaking, their presence keeps the completeness and consistency of the all the tensor operations. For example, given x = mx.nd.array([1, 2, 3]), x[0] should return 1, instead of [1], and x[0:0] should return [] with shape equal to (0,). Zero-size tensors can also be convenient in easing the code logic as placeholders for accumulations or aggregations.

I find this thread provides very good insights on the importance of zero-dim tensors.

How are they implemented?

In the backend (C++), we use ndim = -1 to represent unknown shapes and dim_size = -1 for unknown dim sizes. Before, they were represented by 0.

Is backward compatibility guaranteed in this PR?

Yes, the backward compatibility is guaranteed by default. That means 0 still represents unknown ndim or dim size in any frontend language bindings. It's just converted to -1 in the infer shape logic of the backend.

How to enable zero-dim or zero-size tensors?

Since we are committed to keep backward compatibility, we provided APIs for users to decide whether to opt in for this NumPy compatibility. Users can call mx.set_np_compat(True) to opt in and mx.set_np_compat(False) to opt out. Or in a safer way, use the with statement. To turn on this:

 with mx.np_compat(active=True):
     # A scalar tensor's shape is `()`, whose `ndim` is `0`.
     scalar = mx.nd.ones(shape=())
     assert scalar.shape == ()
 
     # In NumPy compatible mode, 0 in a shape means that dimension contains zero elements.
     data = mx.sym.var("data", shape=(0, 2, 3))
     ret = mx.sym.sin(data)
     arg_shapes, out_shapes, _ = ret.infer_shape()
     assert arg_shapes[0] == (0, 2, 3)
     assert out_shapes[0] == (0, 2, 3)
 
     # -1 means unknown shape dimension size in the new NumPy-compatible shape definition
     data = mx.sym.var("data", shape=(-1, 2, 3))
     ret = mx.sym.sin(data)
     arg_shapes, out_shapes, _ = ret.infer_shape_partial()
     assert arg_shapes[0] == (-1, 2, 3)
     assert out_shapes[0] == (-1, 2, 3)
 
     # When a shape is completely unknown in NumPy-compatible mode, it is
     # represented as `None` in Python.
     data = mx.sym.var("data")
     ret = mx.sym.sin(data)
     arg_shapes, out_shapes, _ = ret.infer_shape_partial()
     assert arg_shapes[0] is None
     assert out_shapes[0] is None

or to disable this:

 with mx.np_compat(active=False):
     # 0 means unknown shape dimension size in the legacy shape definition.
     data = mx.sym.var("data", shape=(0, 2, 3))
     ret = mx.sym.sin(data)
     arg_shapes, out_shapes, _ = ret.infer_shape_partial()
     assert arg_shapes[0] == (0, 2, 3)
     assert out_shapes[0] == (0, 2, 3)
 
     # When a shape is completely unknown in the legacy mode (default), its ndim is
     # equal to 0 and it is represented as `()` in Python.
     data = mx.sym.var("data")
     ret = mx.sym.sin(data)
     arg_shapes, out_shapes, _ = ret.infer_shape_partial()
     assert arg_shapes[0] == ()
     assert out_shapes[0] == ()

Does this mean that every existing operator should support zero-dim or zero-size tensors in this PR?

Please note that the existing operators were implemented when these two types of tensors were not supported in MXNet. Some strong assumptions may have been made in their implementation and hence, lead to errors when NumPy compatibility is turned on. This PR only provides the infrastructure of supporting zero-dim and zero-size tensors in the backend. It does not guarantee that every existing operator would deal with zero-dim/size tensors correctly as in NumPy. As discussed in the RFC, we are going to implement NumPy operators under mxnet.numpy that would fully support these two types of tensors.

[wkcn]

Great work! Thank you!

[eric-haibin-lin]

I reviewed all functions in *.py files, they look good to me

[reminisce]

@hetong007 @yzhliu @sergeykolychev Could you please help review the changes in R, Scala, and Perl respectively? We would like to merge this PR sooner than later since our subsequent work depends on this and it's becoming harder to rebase.

[szha]

@jeremiedb @lanking520 could you help out on this for reviews too? Thanks!

[yzhliu]

LGTM

[lanking520]

Please add some comments on the changes for Scala part.

[jeremiedb]

@hetong007 Were you seing a need for adding integrating a set_np_compat option in the R-package? Otherwise, LGTM.

[TaoLv]

Reviewed changes for mkldnn files. @juliusshufan you might want to have a performance validation on this PR.

[hetong007]

@hetong007 Were you seing a need for adding integrating a set_np_compat option in the R-package? Otherwise, LGTM.

@jeremiedb Would you elaborate more on this need? I think the current R API maps nd.array to native R array so this PR has minimum effect on R-related uses.

[anirudh2290]

Thanks for the decorator! I think customers may not understand the internal workings of custom op and the dispatch to a different thread workflow, and run into this pitfall. Can we synchronize the state of is_numpy_compat when pushing the op to custom op thread. Custom op is one area where customers use numpy operations and I feel there may be users running into this.

[reminisce]

@anirudh2290 That's a good point from users' perspective. But this may not be a trivial implementation by propagating the np_compat from main thread to worker thread as potentially there are other pitfalls. I discussed this @eric-haibin-lin, maybe we can attach a compatibility state to the generated custom op. This can become a follow-up implementation in another PR. This PR does not need to be blocked by this.

[szha]

@anirudh2290 good points on custom op. Given the size of this PR and the fact that users need to opt-in to be affected by this change, I'd recommend we address these points in a follow-up PR, so that @reminisce won't have to keep full-time resolving conflicts :)

[szha]

python, C API, and operator changes LGTM

[anirudh2290]

I am okay with making custom op changes later. just to be sure, the the is_numpy_compat thread local variable is set in the main thread, and is expected to be accessed only from main thread ? For example, this wont cause any issue the multi threaded inference interface provided by frontends like Scala ?

addressed

[junrushao]

LGTM

[reminisce]

@anirudh2290 Thanks for the review. In practice, np_compat needs to be set in the same thread as the one invoking shape inference functions, because that's where zero-dim and zero-size shapes are treated as unknown. In your example, the custom op's forward function invokes mx.nd.ones whose shape inference function was executed on the worker thread, and hence leads to failure. We are going to move all ops towards the direction of being numpy compatible, and that should include the custom op, where the forward/backward functions should always be scoped by the numpy-compatible state, and this problem can be resolved without asking users to use the decorator.

I'm not familiar with how scala multi-threaded inference is implemented, but as long as the np_compat is set in the same thread as the one invoking shape inference functions, the result should be as expected.

[TaoLv]

Verified performance and accuracy for several CNN models with MKL-DNN backend. They look good to me.

[leezu]

@reminisce is this #if 0 intentional? Will it be removed eventually?