dmlc--dgl
Minjie Wang
5dd35580f7
* improve performance of sample_neighbors * some more improve * test script * benchmarks * multi process * update more tests * WIP * adding two API for state saving * add create from state * upd test * missing file * wip: pickle/unpickle * more c apis * find the problem of empty data array * add null array; pickling speed is bad * still bad perf * still bad perf * wip * fix the pickle speed test; now everything looks good * minor fix * bugfix * some lint fix * address comments * more fix * fix lint * add utest for random.choice * add utest for dgl.rand_graph * fix cpp utests * try fix ci * fix bug in TF backend * upd choice docstring * address comments * upd * try fix compile * add comment
79 行
2.5 KiB
Python
79 行
2.5 KiB
Python
"""Container data structures used in DGL runtime.
|
|
reference: tvm/python/tvm/collections.py
|
|
"""
|
|
from __future__ import absolute_import as _abs
|
|
from ._ffi.object import ObjectBase, register_object
|
|
from . import _api_internal
|
|
|
|
@register_object
|
|
class List(ObjectBase):
|
|
"""List container of DGL.
|
|
|
|
You do not need to create List explicitly.
|
|
Normally python list and tuple will be converted automatically
|
|
to List during dgl function call.
|
|
You may get List in return values of DGL function call.
|
|
"""
|
|
def __getitem__(self, i):
|
|
if isinstance(i, slice):
|
|
start = i.start if i.start is not None else 0
|
|
stop = i.stop if i.stop is not None else len(self)
|
|
step = i.step if i.step is not None else 1
|
|
if start < 0:
|
|
start += len(self)
|
|
if stop < 0:
|
|
stop += len(self)
|
|
return [self[idx] for idx in range(start, stop, step)]
|
|
|
|
if i < -len(self) or i >= len(self):
|
|
raise IndexError("List index out of range. List size: {}, got index {}"
|
|
.format(len(self), i))
|
|
if i < 0:
|
|
i += len(self)
|
|
return _api_internal._ListGetItem(self, i)
|
|
|
|
def __len__(self):
|
|
return _api_internal._ListSize(self)
|
|
|
|
@register_object
|
|
class Map(ObjectBase):
|
|
"""Map container of DGL.
|
|
|
|
You do not need to create Map explicitly.
|
|
Normally python dict will be converted automaticall to Map during dgl function call.
|
|
You can use convert to create a dict[ObjectBase-> ObjectBase] into a Map
|
|
"""
|
|
def __getitem__(self, k):
|
|
return _api_internal._MapGetItem(self, k)
|
|
|
|
def __contains__(self, k):
|
|
return _api_internal._MapCount(self, k) != 0
|
|
|
|
def items(self):
|
|
"""Get the items from the map"""
|
|
akvs = _api_internal._MapItems(self)
|
|
return [(akvs[i], akvs[i+1]) for i in range(0, len(akvs), 2)]
|
|
|
|
def __len__(self):
|
|
return _api_internal._MapSize(self)
|
|
|
|
|
|
@register_object
|
|
class StrMap(Map):
|
|
"""A special map container that has str as key.
|
|
|
|
You can use convert to create a dict[str->ObjectBase] into a Map.
|
|
"""
|
|
def items(self):
|
|
"""Get the items from the map"""
|
|
akvs = _api_internal._MapItems(self)
|
|
return [(akvs[i].data, akvs[i+1]) for i in range(0, len(akvs), 2)]
|
|
|
|
@register_object
|
|
class Value(ObjectBase):
|
|
"""Object wrapper for various values."""
|
|
@property
|
|
def data(self):
|
|
"""Return the value data."""
|
|
return _api_internal._ValueGet(self)
|