[MXNet] matplotlib >=3.5 raises TypeError with ax.plot_wireframe in MXNet ndarray

[AnirudhDagar]

With the latest version of matplotlib, multiple notebooks fail with a type error in mxnet (mxnet==1.7.0 & CUDA 10.2). Some of the affected sections include optimization intro, integral calculus, multivariable calculus etc.

TypeError: no implementation found for 'numpy.column_stack' on types that implement __array_function__: [<class 'mxnet.numpy.ndarray'>, <class 'numpy.ndarray'>]

Please see attached traceback and reproduction instructions below.

Steps to reproduce the issue.

1. Setup the d2l environment (using static/build.yml)
2. While setting up the environment, it will automatically install the latest version of matplotlib (i.e. matplotlib==3.5.1 as of today).

Run one of the notebooks which is affected (mentioned above)

Click to expand: Error Traceback

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [7], in <module>
      9 # Plot function
     10 ax = d2l.plt.figure().add_subplot(111, projection='3d')
---> 11 ax.plot_wireframe(x, y, z, **{'rstride': 10, 'cstride': 10})
     12 ax.plot_wireframe(x, y, w, **{'rstride': 10, 'cstride': 10}, color='purple')
     13 d2l.plt.xlabel('x')

File ~/miniconda3/envs/mpl_d2l/lib/python3.8/site-packages/matplotlib/_api/deprecation.py:412, in delete_parameter.<locals>.wrapper(*inner_args, **inner_kwargs)
    402     deprecation_addendum = (
    403         f"If any parameter follows {name!r}, they should be passed as "
    404         f"keyword, not positionally.")
    405     warn_deprecated(
    406         since,
    407         name=repr(name),
   (...)
    410                  else deprecation_addendum,
    411         **kwargs)
--> 412 return func(*inner_args, **inner_kwargs)

File ~/miniconda3/envs/mpl_d2l/lib/python3.8/site-packages/mpl_toolkits/mplot3d/axes3d.py:1908, in Axes3D.plot_wireframe(self, X, Y, Z, *args, **kwargs)
   1906 linec = art3d.Line3DCollection(lines, *args, **kwargs)
   1907 self.add_collection(linec)
-> 1908 self.auto_scale_xyz(X, Y, Z, had_data)
   1910 return linec

File ~/miniconda3/envs/mpl_d2l/lib/python3.8/site-packages/mpl_toolkits/mplot3d/axes3d.py:658, in Axes3D.auto_scale_xyz(self, X, Y, Z, had_data)
    656     self.xy_dataLim.update_from_data_y(Y, not had_data)
    657 if Z is not None:
--> 658     self.zz_dataLim.update_from_data_x(Z, not had_data)
    659 # Let autoscale_view figure out how to use this data.
    660 self.autoscale_view()

File ~/miniconda3/envs/mpl_d2l/lib/python3.8/site-packages/matplotlib/transforms.py:922, in Bbox.update_from_data_x(self, x, ignore)
    906 """
    907 Update the x-bounds of the `Bbox` based on the passed in data. After
    908 updating, the bounds will have positive *width*, and *x0* will be the
   (...)
    919    - When ``None``, use the last value passed to :meth:`ignore`.
    920 """
    921 x = np.ravel(x)
--> 922 self.update_from_data_xy(np.column_stack([x, np.ones(x.size)]),
    923                          ignore=ignore, updatey=False)

File <__array_function__ internals>:180, in column_stack(*args, **kwargs)

TypeError: no implementation found for 'numpy.column_stack' on types that implement __array_function__: [<class 'mxnet.numpy.ndarray'>, <class 'numpy.ndarray'>]

This is another issue validating the need of #2044.

A simple solution for now is to pin the matplotlib version to 1.4. I'll send a PR for this.

cc @astonzhang

[AnirudhDagar]

This is now completely fixed in #2078