dmlc--dgl
xiang song(charlie.song)
a7e941c379
* Add sparse embedding for dgl and update rgcn example * upd * Fix * Revert "Fix" This reverts commit 4da87cdfb8b8c3506b7fc7376cd2385ba8045c2a. * Fix * upd * upd * Fix * Add unitest and update impl * fix * Clean up rgcn example code * upd * upd * update * Fix * update score * sparse for sage * remove model sparse * upd * upd * remove global norm * revert delete model_sparse.py * update according to comments * Fix doc * upd * Fix test * upd * lint * lint * lint * upd * upd * clean up Co-authored-by: Ubuntu <ubuntu@ip-172-31-56-220.ec2.internal>
49 行
1.3 KiB
Python
49 行
1.3 KiB
Python
|
|
"""Shared memory utilities."""
|
|
from .. import backend as F
|
|
from .._ffi.ndarray import empty_shared_mem
|
|
|
|
def get_shared_mem_array(name, shape, dtype):
|
|
""" Get a tensor from shared memory with specific name
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The unique name of the shared memory
|
|
shape : tuple of int
|
|
The shape of the returned tensor
|
|
dtype : F.dtype
|
|
The dtype of the returned tensor
|
|
|
|
Returns
|
|
-------
|
|
F.tensor
|
|
The tensor got from shared memory.
|
|
"""
|
|
name = 'DGL_'+name
|
|
new_arr = empty_shared_mem(name, False, shape, F.reverse_data_type_dict[dtype])
|
|
dlpack = new_arr.to_dlpack()
|
|
return F.zerocopy_from_dlpack(dlpack)
|
|
|
|
def create_shared_mem_array(name, shape, dtype):
|
|
""" Create a tensor from shared memory with the specific name
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The unique name of the shared memory
|
|
shape : tuple of int
|
|
The shape of the returned tensor
|
|
dtype : F.dtype
|
|
The dtype of the returned tensor
|
|
|
|
Returns
|
|
-------
|
|
F.tensor
|
|
The created tensor.
|
|
"""
|
|
name = 'DGL_'+name
|
|
new_arr = empty_shared_mem(name, True, shape, F.reverse_data_type_dict[dtype])
|
|
dlpack = new_arr.to_dlpack()
|
|
return F.zerocopy_from_dlpack(dlpack)
|