Incubator-mxnet: net.optimize_for doesn't work with numpy semantics

[bgawrych]

Description

Forward pass after calling optimize_for with specific backend doesn't work. I'm not sure what this error mean, but found a way to overcome this (ugly way :))
Problem occurs on master and 1.x branches

Error Message

Traceback (most recent call last):
  File "../d.py", line 23, in <module>
    print(net(a, b))
  File "/home/bgawrych/Desktop/mxnet/python/mxnet/gluon/block.py", line 1407, in __call__
    return super().__call__(x, *args)
  File "/home/bgawrych/Desktop/mxnet/python/mxnet/gluon/block.py", line 716, in __call__
    _check_all_np_ndarrays(out)
  File "/home/bgawrych/Desktop/mxnet/python/mxnet/gluon/utils.py", line 480, in _check_all_np_ndarrays
    raise TypeError("Block's output ndarrays/symbols must be of type `mxnet.numpy.ndarray`"
TypeError: Block's output ndarrays/symbols must be of type `mxnet.numpy.ndarray` or `mxnet.symbol.numpy._Symbol`, while got output type <class 'mxnet.ndarray.ndarray.NDArray'>

To Reproduce

import mxnet as mx
from mxnet.gluon import HybridBlock

mx.npx.set_np()

class TestBlock(HybridBlock):
    def __init__(self):
        super(TestBlock, self).__init__()
        self.d = mx.gluon.nn.Dense(1)
    def hybrid_forward(self, F, a, b, *args):
        res = self.d.hybrid_forward(F, a, b)
        return res

a = mx.np.random.uniform(low=-1, high=1, size=(1,1))
b = mx.np.random.uniform(low=-1, high=1, size=(1,1))

net = TestBlock()
net.initialize()
net.hybridize()

print(net(a, b))
net.optimize_for(a, b, backend="MKLDNN")
#print(net(a, b)) # <---- this line doesn't work now - we need to reload symbol with JSON
inputs, sym = net._cached_graph
sym = mx.sym.np._symbol.load_json(sym.tojson())
x = mx.gluon.SymbolBlock(sym, [mx.sym.var('data0'), mx.sym.var('data1')], net.collect_params())

print(x(a, b))

What have you tried to solve it?

1. Add ConvertShapeAttrToNumPyCompatible(&g); in MXOptimizeForBackend- doesn't help

@samskalicky maybe you will be able to help

[samskalicky]

Thanks for the code to reproduce @bgawrych, I can reproduce on teh v1.8.x branch. will take a look

[samskalicky]

After checking with @leezu it looks like we need to do something like:

# Partition the graph.                                                                                                               
out = out.optimize_for(self._backend, arg_dict, aux_dict, ctx, **self._backend_opts)

# convert to numpy symbol if needed                                                                                                  
if _mx_npx.is_np_array():
    out = out.as_np_ndarray()

#update cached graph with partitioned graph                                                                                          
self._cached_graph = data, out

here where we call optimize_for in the Gluon block:
https://github.com/apache/incubator-mxnet/blob/0faecf01eddc781793927cb7cb9415310e9b21a2/python/mxnet/gluon/block.py#L1039
With this change it seems to be working. I'll create a PR with this on v1.x and master branches. I guess there hasnt been any testing of Numpy functionality with optimize_for yet....