GitLab

[Feature] Add `state_dict`, `load_state_dict`, `param_groups` to `dgl.optim.SparseGradOptimizer` (#5311)

* init update

* all get/set optm_state

* add unit tests

* add docstring

* fix for multiple embeddings

* move embedding methods to private

* fix lint

* fix unit tests

* resolve comments

* merge master

* expose GeneratePermutation

* add sparse_all_to_all_push

* add sparse_all_to_all_pull

* add unit test

* handle world_size=1

* remove python nccl wrapper

* remove the nccl dependency

* use pinned memory to speedup D2H copy

* fix lint

* resolve comments

* fix lint

* fix ut

* resolve comments

* fix duplicate worker_init_fn

* lint

* lint again

* uugh

* [Sparse] Support converson to/from torch sparse tensor.

* Update

* [Sparse] Support column-wise softmax

* Update python/dgl/sparse/softmax.py
Co-authored-by: Mufei Li <mufeili1996@gmail.com>

---------
Co-authored-by: Mufei Li <mufeili1996@gmail.com>

* subgraphx commit

* nits

* newline eof added

* lint fix

* test script updated to use default values

* lint fix

* graphs that are used for test cases are updated to a small graph

* lint formatted

* test paramter adj to complete the test under 20s

* lint fixes

---------
Co-authored-by: kxm180046 <kxm180046@utdallas.edu>

* [Misc] Auto-format tests.

* more

---------
Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>

* rename

* remove test transpose

* udpate

* rename

* cleanup

---------
Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>