Update Parametric UMAP to support torch backend.

[timsainb]

There still appear to be some issues when using a jax backend (nans in network weights) that I haven't been able to debug but everything else seems to work.

[pep8speaks]

Hello @timsainb! Thanks for opening this PR. We checked the lines you've touched for PEP 8 issues, and found:

• In the file umap/parametric_umap.py:

Line 26:53: W291 trailing whitespace
Line 37:89: E501 line too long (117 > 88 characters)
Line 49:26: E251 unexpected spaces around keyword / parameter equals
Line 49:28: E251 unexpected spaces around keyword / parameter equals
Line 52:22: E251 unexpected spaces around keyword / parameter equals
Line 52:24: E251 unexpected spaces around keyword / parameter equals
Line 80:89: E501 line too long (116 > 88 characters)
Line 81:89: E501 line too long (103 > 88 characters)
Line 82:89: E501 line too long (120 > 88 characters)
Line 82:121: W291 trailing whitespace
Line 83:89: E501 line too long (93 > 88 characters)
Line 126:9: E265 block comment should start with '# '
Line 215:26: E251 unexpected spaces around keyword / parameter equals
Line 215:28: E251 unexpected spaces around keyword / parameter equals
Line 277:13: E265 block comment should start with '# '
Line 277:48: W291 trailing whitespace
Line 602:1: E302 expected 2 blank lines, found 1
Line 723:1: W293 blank line contains whitespace
Line 926:22: E251 unexpected spaces around keyword / parameter equals
Line 926:24: E251 unexpected spaces around keyword / parameter equals
Line 947:17: W291 trailing whitespace
Line 949:1: W293 blank line contains whitespace
Line 951:89: E501 line too long (100 > 88 characters)
Line 954:89: E501 line too long (103 > 88 characters)
Line 955:1: W293 blank line contains whitespace
Line 962:1: W293 blank line contains whitespace
Line 965:1: W293 blank line contains whitespace
Line 966:1: W293 blank line contains whitespace
Line 986:1: W293 blank line contains whitespace
Line 994:1: W293 blank line contains whitespace