Drop MXNet Support

[matthewfeickert]

Description

MXNet isn't going to be implementing statistical distributions anytime soon, so it seems like they will be diverging from the rest of our ML framework backends even more. Instead of keeping MXNet around as a vestigial dependency I would suggest we just drop support for it.

Thoughts @kratsg @lukasheinrich?

[kratsg]

MXNet isn't going to be implementing statistical distributions anytime soon

how do we know?

[matthewfeickert]

how do we know?

This issue makes it seem like having dedicated distributions with full functionality like TF and PyTorch isn't going to happen right away. The PR that is linked at the very end of it does add sampling (presumably in the next release) but it again lacks the full feature set of the other backends.

In my mind it makes more sense to just discontinue any effort on MXNet and think more about PR #377 instead.

[lukasheinrich]

I think it would be a good idea to replace MXNet with jax (#377 ) it's a more fully featured backend that has

• hardware support
• autodifferentiation
• very similar to numpy

and could actually be the one backend that makes things faster than numpy (currently pytorch and tf are only faster in very extreme cases)

[kratsg]

Does jax support distributions?

[lukasheinrich]

yeah: https://jax.readthedocs.io/en/latest/jax.scipy.html#module-jax.scipy.stats

[matthewfeickert]

@kratsg Given the 👍 I assume you have no qualms with me opening a MR for a branch in which I've already done this?

[kratsg]

yes, no qualms