Improve createOpState for custom ops

[samskalicky]

Description

Support passing all args from createOpState to custom ops: context, input shapes & types

• Also updates test_gemm.py and test_relu.py for custom ops on 2.0 now that Symbol.bind is deprecated

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

Changes

• Feature1, tests, (and when applicable, API doc)
• Feature2, tests, (and when applicable, API doc)

Comments

• If this change is a backward incompatible change, why must this change be made.
• Interesting edge cases to note here

[mxnet-bot]

Hey @samskalicky , Thanks for submitting the PR
All tests are already queued to run once. If tests fail, you can trigger one or more tests again with the following commands:

• To trigger all jobs: @mxnet-bot run ci [all]
• To trigger specific jobs: @mxnet-bot run ci [job1, job2]

---

CI supported jobs: [unix-gpu, miscellaneous, edge, sanity, windows-gpu, unix-cpu, centos-gpu, clang, centos-cpu, windows-cpu, website]

---

Note:
Only following 3 categories can trigger CI :PR Author, MXNet Committer, Jenkins Admin.
All CI tests must pass before the PR can be merged.

[samskalicky]

@mseth10 @rondogency this is a small PR with just the change for custom stateful ops. Please review, thanks!

[szha]

thanks for the enhancement! as we are approaching the (beta) release of mxnet, do you think the API would be finalized with this PR?

[samskalicky]

I think so, with this PR the examples are all up to date now with Gluon (having removed the symbol bind/executor stuff).

As with anything, as it gets more use we'll find ways to improve things. But we're ready to start helping users extend MXNet!

[mseth10]

LGTM

[karan6181]

LGTM. Thank You for you contribution!

[rondogency]

One quick request, in terms of the test, can you add an end-to-end test training example written in gluon like this one: #18167, so that we can have better test suite on performance to show to user

[rondogency]

LGTM! the above comment can be addressed in the future PR