Adapted use of mx module. Applied isort.

[jaheba]

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[codecov-commenter]

Codecov Report

Merging #878 into master will increase coverage by 48.08%.
The diff coverage is 99.12%.

@@             Coverage Diff             @@
##           master     #878       +/-   ##
===========================================
+ Coverage   37.73%   85.82%   +48.08%     
===========================================
  Files         195      195               
  Lines       11786    11783        -3     
===========================================
+ Hits         4448    10113     +5665     
+ Misses       7338     1670     -5668     

Impacted Files	Coverage Δ
src/gluonts/block.py	0.00% <0.00%> (-100.00%)	⬇️
src/gluonts/dataset/artificial/__init__.py	80.00% <ø> (ø)
src/gluonts/dataset/common.py	92.78% <ø> (+47.11%)	⬆️
src/gluonts/model/tpp/predictor.py	85.41% <ø> (+39.58%)	⬆️
src/gluonts/model/transformer/trans_decoder.py	100.00% <ø> (+69.69%)	⬆️
src/gluonts/model/transformer/trans_encoder.py	100.00% <ø> (+62.50%)	⬆️
src/gluonts/nursery/sagemaker_sdk/defaults.py	100.00% <ø> (ø)
src/gluonts/runtime_params.py	100.00% <ø> (ø)
src/gluonts/shell/serve/__init__.py	90.00% <ø> (+40.00%)	⬆️
src/gluonts/support/util.py	94.23% <ø> (+65.38%)	⬆️
... and 267 more