Skip to content

GitLab

Unverified Commit 44089c8b authored by Minjie Wang's avatar Minjie Wang Committed by GitHub
Browse files

[Refactor][Graph] Merge DGLGraph and DGLHeteroGraph (#1862) 



* Merge

* [Graph][CUDA] Graph on GPU and many refactoring (#1791)

* change edge_ids behavior and C++ impl

* fix unittests; remove utils.Index in edge_id

* pass mx and th tests

* pass tf test

* add aten::Scatter_

* Add nonzero; impl CSRGetDataAndIndices/CSRSliceMatrix

* CSRGetData and CSRGetDataAndIndices passed tests

* CSRSliceMatrix basic tests

* fix bug in empty slice

* CUDA CSRHasDuplicate

* has_node; has_edge_between

* predecessors, successors

* deprecate send/recv; fix send_and_recv

* deprecate send/recv; fix send_and_recv

* in_edges; out_edges; all_edges; apply_edges

* in deg/out deg

* subgraph/edge_subgraph

* adj

* in_subgraph/out_subgraph

* sample neighbors

* set/get_n/e_repr

* wip: working on refactoring all idtypes

* pass ndata/edata tests on gpu

* fix

* stash

* workaround nonzero issue

* stash

* nx conversion

* test_hetero_basics except update routines

* test_update_routines

* test_hetero_basics for pytorch

* more fixes

* WIP: flatten graph

* wip: flatten

* test_flatten

* test_to_device

* fix bug in to_homo

* fix bug in CSRSliceMatrix

* pass subgraph test

* fix send_and_recv

* fix filter

* test_heterograph

* passed all pytorch tests

* fix mx unittest

* fix pytorch test_nn

* fix all unittests for PyTorch

* passed all mxnet tests

* lint

* fix tf nn test

* pass all tf tests

* lint

* lint

* change deprecation

* try fix compile

* lint

* update METIDS

* fix utest

* fix

* fix utests

* try debug

* revert

* small fix

* fix utests

* upd

* upd

* upd

* fix

* upd

* upd

* upd

* upd

* upd

* trigger

* +1s

* [kernel] Use heterograph index instead of unitgraph index (#1813)

* upd

* upd

* upd

* fix

* upd

* upd

* upd

* upd

* upd

* trigger

* +1s

* [Graph] Mutation for Heterograph (#1818)

* mutation add_nodes and add_edges

* Add support for remove_edges, remove_nodes, add_selfloop, remove_selfloop

* Fix
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-51-214.ec2.internal>

* upd

* upd

* upd

* fix

* [Transfom] Mutable transform (#1833)

* add nodesy

* All three

* Fix

* lint

* Add some test case

* Fix

* Fix

* Fix

* Fix

* Fix

* Fix

* fix

* triger

* Fix

* fix
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-51-214.ec2.internal>

* [Graph] Migrate Batch & Readout module to heterograph (#1836)

* dgl.batch

* unbatch

* fix to device

* reduce readout; segment reduce

* change batch_num_nodes|edges to function

* reduce readout/ softmax

* broadcast

* topk

* fix

* fix tf and mx

* fix some ci

* fix batch but unbatch differently

* new checkk

* upd

* upd

* upd

* idtype behavior; code reorg

* idtype behavior; code reorg

* wip: test_basics

* pass test_basics

* WIP: from nx/ to nx

* missing files

* upd

* pass test_basics:test_nx_conversion

* Fix test

* Fix inplace update

* WIP: fixing tests

* upd

* pass test_transform cpu

* pass gpu test_transform

* pass test_batched_graph

* GPU graph auto cast to int32

* missing file

* stash

* WIP: rgcn-hetero

* Fix two datasety

* upd

* weird

* Fix capsuley

* fuck you

* fuck matthias

* Fix dgmg

* fix bug in block degrees; pass rgcn-hetero

* rgcn

* gat and diffpool fix
also fix ppi and tu dataset

* Tree LSTM

* pointcloud

* rrn; wip: sgc

* resolve conflicts

* upd

* sgc and reddit dataset

* upd

* Fix deepwalk, gindt and gcn

* fix datasets and sign

* optimization

* optimization

* upd

* upd

* Fix GIN

* fix bug in add_nodes add_edges; tagcn

* adaptive sampling and gcmc

* upd

* upd

* fix geometric

* fix

* metapath2vec

* fix agnn

* fix pickling problem of block

* fix utests

* miss file

* linegraph

* upd

* upd

* upd

* graphsage

* stgcn_wave

* fix hgt

* on unittests

* Fix transformer

* Fix HAN

* passed pytorch unittests

* lint

* fix

* Fix cluster gcn

* cluster-gcn is ready

* on fixing block related codes

* 2nd order derivative

* Revert "2nd order derivative"

This reverts commit 523bf6c249bee61b51b1ad1babf42aad4167f206.

* passed torch utests again

* fix all mxnet unittests

* delete some useless tests

* pass all tf cpu tests

* disable

* disable distributed unittest

* fix

* fix

* lint

* fix

* fix

* fix script

* fix tutorial

* fix apply edges bug

* fix 2 basics

* fix tutorial
Co-authored-by: default avataryzh119 <expye@outlook.com>
Co-authored-by: default avatarxiang song(charlie.song) <classicxsong@gmail.com>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-51-214.ec2.internal>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-7-42.us-west-2.compute.internal>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-1-5.us-west-2.compute.internal>
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-68-185.ec2.internal>
parent 015acfd2
Hide whitespace changes
Inline Side-by-side
Showing with 1460 additions and 1925 deletions
src/graph/pickle.cc
View file @ 44089c8b
......@@ -10,6 +10,7 @@
#include <dmlc/memory_io.h>
#include "./heterograph.h"
#include "../c_api_common.h"
#include "unit_graph.h"
using namespace dgl::runtime;
......@@ -22,7 +23,7 @@ HeteroPickleStates HeteroPickle(HeteroGraphPtr graph) {
strm->Write(ImmutableGraph::ToImmutable(graph->meta_graph()));
strm->Write(graph->NumVerticesPerType());
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
SparseFormat fmt = graph->SelectFormat(etype, SparseFormat::kAny);
SparseFormat fmt = graph->SelectFormat(etype, all_code);
switch (fmt) {
case SparseFormat::kCOO: {
strm->Write(SparseFormat::kCOO);
......@@ -82,7 +83,8 @@ HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states) {
CHECK(strm->Read(&rsorted)) << "Invalid flag 'rsorted'";
CHECK(strm->Read(&csorted)) << "Invalid flag 'csorted'";
auto coo = aten::COOMatrix(num_src, num_dst, row, col, aten::NullArray(), rsorted, csorted);
relgraph = CreateFromCOO(num_vtypes, coo);
// TODO(zihao) fix
relgraph = CreateFromCOO(num_vtypes, coo, all_code);
break;
}
case SparseFormat::kCSR: {
......@@ -93,7 +95,8 @@ HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states) {
bool sorted;
CHECK(strm->Read(&sorted)) << "Invalid flag 'sorted'";
auto csr = aten::CSRMatrix(num_src, num_dst, indptr, indices, edge_id, sorted);
relgraph = CreateFromCSR(num_vtypes, csr);
// TODO(zihao) fix
relgraph = CreateFromCSR(num_vtypes, csr, all_code);
break;
}
case SparseFormat::kCSC:
......
src/graph/sampling/neighbor/neighbor.cc
View file @ 44089c8b
......@@ -62,7 +62,7 @@ HeteroSubgraph SampleNeighbors(
induced_edges[etype] = earr.id;
} else {
// sample from one relation graph
auto req_fmt = (dir == EdgeDir::kOut)? SparseFormat::kCSR : SparseFormat::kCSC;
auto req_fmt = (dir == EdgeDir::kOut)? csr_code : csc_code;
auto avail_fmt = hg->SelectFormat(etype, req_fmt);
COOMatrix sampled_coo;
switch (avail_fmt) {
......@@ -148,7 +148,7 @@ HeteroSubgraph SampleNeighborsTopk(
induced_edges[etype] = earr.id;
} else {
// sample from one relation graph
auto req_fmt = (dir == EdgeDir::kOut)? SparseFormat::kCSR : SparseFormat::kCSC;
auto req_fmt = (dir == EdgeDir::kOut)? csr_code : csc_code;
auto avail_fmt = hg->SelectFormat(etype, req_fmt);
COOMatrix sampled_coo;
switch (avail_fmt) {
......
src/graph/serialize/dglgraph_data.h
View file @ 44089c8b
......@@ -20,6 +20,7 @@
#include <vector>
#include <algorithm>
#include <utility>
#include <memory>
#include "../../c_api_common.h"
using dgl::runtime::NDArray;
......
src/graph/serialize/heterograph_data.h
View file @ 44089c8b
......@@ -20,6 +20,7 @@
#include <string>
#include <utility>
#include <vector>
#include <memory>
#include "../../c_api_common.h"
#include "../heterograph.h"
......
src/graph/serialize/streamwithcount.h
View file @ 44089c8b
......@@ -8,6 +8,7 @@
#include <dmlc/io.h>
#include <dmlc/type_traits.h>
#include <memory>
/*!
* \brief StreamWithCount counts the bytes that already written into the
......
src/graph/transform/remove_edges.cc
View file @ 44089c8b
......@@ -28,7 +28,7 @@ RemoveEdges(const HeteroGraphPtr graph, const std::vector<IdArray> &eids) {
const int64_t num_etypes = graph->NumEdgeTypes();
for (int64_t etype = 0; etype < num_etypes; ++etype) {
const SparseFormat fmt = graph->SelectFormat(etype, SparseFormat::kCOO);
const SparseFormat fmt = graph->SelectFormat(etype, coo_code);
const auto src_dst_types = graph->GetEndpointTypes(etype);
const dgl_type_t srctype = src_dst_types.first;
const dgl_type_t dsttype = src_dst_types.second;
......
src/graph/transform/union_partition.cc
View file @ 44089c8b
......@@ -24,9 +24,9 @@ HeteroGraphPtr JointUnionHeteroGraph(
HeteroGraphPtr rgptr = nullptr;
// ALL = CSC | CSR | COO
dgl_format_code_t format = (1 << (SparseFormat2Code(SparseFormat::kCOO)-1)) |
(1 << (SparseFormat2Code(SparseFormat::kCSR)-1)) |
(1 << (SparseFormat2Code(SparseFormat::kCSC)-1));
const dgl_format_code_t code =\
component_graphs[0]->GetRelationGraph(etype)->GetAllowedFormats();
// get common format
for (size_t i = 0; i < component_graphs.size(); ++i) {
const auto& cg = component_graphs[i];
......@@ -35,19 +35,13 @@ HeteroGraphPtr JointUnionHeteroGraph(
CHECK_EQ(num_dst_v, component_graphs[i]->NumVertices(dst_vtype)) << "Input graph[" << i <<
"] should have same number of dst vertices as input graph[0]";
const std::string restrict_format = cg->GetRelationGraph(etype)->GetRestrictFormat();
const SparseFormat curr_format = ParseSparseFormat(restrict_format);
if (curr_format == SparseFormat::kCOO ||
curr_format == SparseFormat::kCSR ||
curr_format == SparseFormat::kCSC)
format &=(1 << (SparseFormat2Code(curr_format)-1));
const dgl_format_code_t curr_code = cg->GetRelationGraph(etype)->GetAllowedFormats();
if (curr_code != code)
LOG(FATAL) << "All components should have the same formats";
}
CHECK_GT(format, 0) << "The conjunction of restrict_format of the relation graphs under " <<
etype << "should not be None.";
// prefer COO
if (FORMAT_HAS_COO(format)) {
if (FORMAT_HAS_COO(code)) {
std::vector<aten::COOMatrix> coos;
for (size_t i = 0; i < component_graphs.size(); ++i) {
const auto& cg = component_graphs[i];
......@@ -57,9 +51,8 @@ HeteroGraphPtr JointUnionHeteroGraph(
aten::COOMatrix res = aten::UnionCoo(coos);
rgptr = UnitGraph::CreateFromCOO(
(src_vtype == dst_vtype) ? 1 : 2, res,
SparseFormat::kAny);
} else if (FORMAT_HAS_CSR(format)) {
(src_vtype == dst_vtype) ? 1 : 2, res, code);
} else if (FORMAT_HAS_CSR(code)) {
std::vector<aten::CSRMatrix> csrs;
for (size_t i = 0; i < component_graphs.size(); ++i) {
const auto& cg = component_graphs[i];
......@@ -69,9 +62,8 @@ HeteroGraphPtr JointUnionHeteroGraph(
aten::CSRMatrix res = aten::UnionCsr(csrs);
rgptr = UnitGraph::CreateFromCSR(
(src_vtype == dst_vtype) ? 1 : 2, res,
SparseFormat::kAny);
} else if (FORMAT_HAS_CSC(format)) {
(src_vtype == dst_vtype) ? 1 : 2, res, code);
} else if (FORMAT_HAS_CSC(code)) {
// CSR and CSC have the same storage format, i.e. CSRMatrix
std::vector<aten::CSRMatrix> cscs;
for (size_t i = 0; i < component_graphs.size(); ++i) {
......@@ -82,8 +74,7 @@ HeteroGraphPtr JointUnionHeteroGraph(
aten::CSRMatrix res = aten::UnionCsr(cscs);
rgptr = UnitGraph::CreateFromCSC(
(src_vtype == dst_vtype) ? 1 : 2, res,
SparseFormat::kAny);
(src_vtype == dst_vtype) ? 1 : 2, res, code);
}
rel_graphs[etype] = rgptr;
......@@ -108,30 +99,22 @@ HeteroGraphPtr DisjointUnionHeteroGraph2(
uint64_t src_offset = 0, dst_offset = 0;
HeteroGraphPtr rgptr = nullptr;
// ALL = CSC | CSR | COO
dgl_format_code_t format = (1 << (SparseFormat2Code(SparseFormat::kCOO)-1)) |
(1 << (SparseFormat2Code(SparseFormat::kCSR)-1)) |
(1 << (SparseFormat2Code(SparseFormat::kCSC)-1));
const dgl_format_code_t code =\
component_graphs[0]->GetRelationGraph(etype)->GetAllowedFormats();
// do some preprocess
for (size_t i = 0; i < component_graphs.size(); ++i) {
const auto& cg = component_graphs[i];
const std::string restrict_format = cg->GetRelationGraph(etype)->GetRestrictFormat();
const SparseFormat curr_format = ParseSparseFormat(restrict_format);
if (curr_format == SparseFormat::kCOO ||
curr_format == SparseFormat::kCSR ||
curr_format == SparseFormat::kCSC)
format &=(1 << (SparseFormat2Code(curr_format)-1));
const dgl_format_code_t cur_code = cg->GetRelationGraph(etype)->GetAllowedFormats();
if (cur_code != code)
LOG(FATAL) << "All components should have the same formats";
// Update offsets
src_offset += cg->NumVertices(src_vtype);
dst_offset += cg->NumVertices(dst_vtype);
}
CHECK_GT(format, 0) << "The conjunction of restrict_format of the relation graphs under " <<
etype << "should not be None.";
// prefer COO
if (FORMAT_HAS_COO(format)) {
if (FORMAT_HAS_COO(code)) {
std::vector<aten::COOMatrix> coos;
for (size_t i = 0; i < component_graphs.size(); ++i) {
const auto& cg = component_graphs[i];
......@@ -142,9 +125,8 @@ HeteroGraphPtr DisjointUnionHeteroGraph2(
aten::COOMatrix res = aten::DisjointUnionCoo(coos);
rgptr = UnitGraph::CreateFromCOO(
(src_vtype == dst_vtype) ? 1 : 2, res,
SparseFormat::kAny);
} else if (FORMAT_HAS_CSR(format)) {
(src_vtype == dst_vtype) ? 1 : 2, res, code);
} else if (FORMAT_HAS_CSR(code)) {
std::vector<aten::CSRMatrix> csrs;
for (size_t i = 0; i < component_graphs.size(); ++i) {
const auto& cg = component_graphs[i];
......@@ -155,9 +137,8 @@ HeteroGraphPtr DisjointUnionHeteroGraph2(
aten::CSRMatrix res = aten::DisjointUnionCsr(csrs);
rgptr = UnitGraph::CreateFromCSR(
(src_vtype == dst_vtype) ? 1 : 2, res,
SparseFormat::kAny);
} else if (FORMAT_HAS_CSC(format)) {
(src_vtype == dst_vtype) ? 1 : 2, res, code);
} else if (FORMAT_HAS_CSC(code)) {
// CSR and CSC have the same storage format, i.e. CSRMatrix
std::vector<aten::CSRMatrix> cscs;
for (size_t i = 0; i < component_graphs.size(); ++i) {
......@@ -168,8 +149,7 @@ HeteroGraphPtr DisjointUnionHeteroGraph2(
aten::CSRMatrix res = aten::DisjointUnionCsr(cscs);
rgptr = UnitGraph::CreateFromCSC(
(src_vtype == dst_vtype) ? 1 : 2, res,
SparseFormat::kAny);
(src_vtype == dst_vtype) ? 1 : 2, res, code);
}
rel_graphs[etype] = rgptr;
num_nodes_per_type[src_vtype] = src_offset;
......@@ -226,10 +206,9 @@ std::vector<HeteroGraphPtr> DisjointPartitionHeteroBySizes2(
std::vector<std::vector<HeteroGraphPtr>> rel_graphs;
rel_graphs.resize(batch_size);
// Loop over all edge types
auto format = batched_graph->GetRelationGraph(0)->GetFormatInUse();
auto restrict_format = batched_graph->GetRelationGraph(0)->GetRestrictFormat();
auto code = batched_graph->GetRelationGraph(0)->GetAllowedFormats();
if (FORMAT_HAS_COO(format)) {
if (FORMAT_HAS_COO(code)) {
for (uint64_t etype = 0; etype < num_edge_types; ++etype) {
auto pair = meta_graph->FindEdge(etype);
const dgl_type_t src_vtype = pair.first;
......@@ -242,13 +221,11 @@ std::vector<HeteroGraphPtr> DisjointPartitionHeteroBySizes2(
vertex_cumsum[dst_vtype]);
for (uint64_t g = 0; g < batch_size; ++g) {
HeteroGraphPtr rgptr = UnitGraph::CreateFromCOO(
(src_vtype == dst_vtype) ? 1 : 2, res[g],
ParseSparseFormat(restrict_format));
(src_vtype == dst_vtype) ? 1 : 2, res[g], code);
rel_graphs[g].push_back(rgptr);
}
}
} else if (FORMAT_HAS_CSR(format)) {
} else if (FORMAT_HAS_CSR(code)) {
for (uint64_t etype = 0; etype < num_edge_types; ++etype) {
auto pair = meta_graph->FindEdge(etype);
const dgl_type_t src_vtype = pair.first;
......@@ -261,13 +238,11 @@ std::vector<HeteroGraphPtr> DisjointPartitionHeteroBySizes2(
vertex_cumsum[dst_vtype]);
for (uint64_t g = 0; g < batch_size; ++g) {
HeteroGraphPtr rgptr = UnitGraph::CreateFromCSR(
(src_vtype == dst_vtype) ? 1 : 2, res[g],
ParseSparseFormat(restrict_format));
(src_vtype == dst_vtype) ? 1 : 2, res[g], code);
rel_graphs[g].push_back(rgptr);
}
}
} else if (FORMAT_HAS_CSC(format)) {
} else if (FORMAT_HAS_CSC(code)) {
for (uint64_t etype = 0; etype < num_edge_types; ++etype) {
auto pair = meta_graph->FindEdge(etype);
const dgl_type_t src_vtype = pair.first;
......@@ -281,9 +256,7 @@ std::vector<HeteroGraphPtr> DisjointPartitionHeteroBySizes2(
vertex_cumsum[src_vtype]);
for (uint64_t g = 0; g < batch_size; ++g) {
HeteroGraphPtr rgptr = UnitGraph::CreateFromCSC(
(src_vtype == dst_vtype) ? 1 : 2, res[g],
SparseFormat::kAny);
(src_vtype == dst_vtype) ? 1 : 2, res[g], code);
rel_graphs[g].push_back(rgptr);
}
}
......
src/graph/unit_graph.cc
View file @ 44089c8b
......@@ -209,16 +209,20 @@ class UnitGraph::COO : public BaseHeteroGraph {
IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
return aten::COOGetData(adj_, src, dst);
return aten::COOGetAllData(adj_, src, dst);
}
EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override {
EdgeArray EdgeIdsAll(dgl_type_t etype, IdArray src, IdArray dst) const override {
CHECK(aten::IsValidIdArray(src)) << "Invalid vertex id array.";
CHECK(aten::IsValidIdArray(dst)) << "Invalid vertex id array.";
const auto& arrs = aten::COOGetDataAndIndices(adj_, src, dst);
return EdgeArray{arrs[0], arrs[1], arrs[2]};
}
IdArray EdgeIdsOne(dgl_type_t etype, IdArray src, IdArray dst) const override {
return aten::COOGetData(adj_, src, dst);
}
std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
CHECK(eid < NumEdges(etype)) << "Invalid edge id: " << eid;
const dgl_id_t src = aten::IndexSelect<int64_t>(adj_.row, eid);
......@@ -334,17 +338,17 @@ class UnitGraph::COO : public BaseHeteroGraph {
return aten::CSRMatrix();
}
SparseFormat SelectFormat(dgl_type_t etype, SparseFormat preferred_format) const override {
SparseFormat SelectFormat(dgl_type_t etype, dgl_format_code_t preferred_formats) const override {
LOG(FATAL) << "Not enabled for COO graph";
return SparseFormat::kAny;
return SparseFormat::kCOO;
}
std::string GetRestrictFormat() const override {
dgl_format_code_t GetAllowedFormats() const override {
LOG(FATAL) << "Not enabled for COO graph";
return std::string("");
return 0;
}
dgl_format_code_t GetFormatInUse() const override {
dgl_format_code_t GetCreatedFormats() const override {
LOG(FATAL) << "Not enabled for COO graph";
return 0;
}
......@@ -392,7 +396,7 @@ class UnitGraph::COO : public BaseHeteroGraph {
return subg;
}
HeteroGraphPtr GetGraphInFormat(SparseFormat restrict_format) const override {
HeteroGraphPtr GetGraphInFormat(dgl_format_code_t formats) const override {
LOG(FATAL) << "Not enabled for COO graph.";
return nullptr;
}
......@@ -591,16 +595,20 @@ class UnitGraph::CSR : public BaseHeteroGraph {
IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
return aten::CSRGetData(adj_, src, dst);
return aten::CSRGetAllData(adj_, src, dst);
}
EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override {
EdgeArray EdgeIdsAll(dgl_type_t etype, IdArray src, IdArray dst) const override {
CHECK(aten::IsValidIdArray(src)) << "Invalid vertex id array.";
CHECK(aten::IsValidIdArray(dst)) << "Invalid vertex id array.";
const auto& arrs = aten::CSRGetDataAndIndices(adj_, src, dst);
return EdgeArray{arrs[0], arrs[1], arrs[2]};
}
IdArray EdgeIdsOne(dgl_type_t etype, IdArray src, IdArray dst) const override {
return aten::CSRGetData(adj_, src, dst);
}
std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
LOG(FATAL) << "Not enabled for CSR graph.";
return {};
......@@ -729,17 +737,17 @@ class UnitGraph::CSR : public BaseHeteroGraph {
return adj_;
}
SparseFormat SelectFormat(dgl_type_t etype, SparseFormat preferred_format) const override {
SparseFormat SelectFormat(dgl_type_t etype, dgl_format_code_t preferred_formats) const override {
LOG(FATAL) << "Not enabled for CSR graph";
return SparseFormat::kAny;
return SparseFormat::kCSR;
}
std::string GetRestrictFormat() const override {
LOG(FATAL) << "Not enabled for CSR graph";
return std::string("");
dgl_format_code_t GetAllowedFormats() const override {
LOG(FATAL) << "Not enabled for COO graph";
return 0;
}
dgl_format_code_t GetFormatInUse() const override {
dgl_format_code_t GetCreatedFormats() const override {
LOG(FATAL) << "Not enabled for CSR graph";
return 0;
}
......@@ -765,7 +773,7 @@ class UnitGraph::CSR : public BaseHeteroGraph {
return {};
}
HeteroGraphPtr GetGraphInFormat(SparseFormat restrict_format) const override {
HeteroGraphPtr GetGraphInFormat(dgl_format_code_t formats) const override {
LOG(FATAL) << "Not enabled for CSR graph.";
return nullptr;
}
......@@ -813,11 +821,13 @@ uint8_t UnitGraph::NumBits() const {
}
bool UnitGraph::IsMultigraph() const {
return GetAny()->IsMultigraph();
const SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
return ptr->IsMultigraph();
}
uint64_t UnitGraph::NumVertices(dgl_type_t vtype) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
const SparseFormat fmt = SelectFormat(all_code);
const auto ptr = GetFormat(fmt);
// TODO(BarclayII): we have a lot of special handling for CSC.
// Need to have a UnitGraph::CSC backend instead.
......@@ -831,7 +841,7 @@ uint64_t UnitGraph::NumEdges(dgl_type_t etype) const {
}
bool UnitGraph::HasVertex(dgl_type_t vtype, dgl_id_t vid) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
const SparseFormat fmt = SelectFormat(all_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
vtype = (vtype == SrcType()) ? DstType() : SrcType();
......@@ -844,7 +854,7 @@ BoolArray UnitGraph::HasVertices(dgl_type_t vtype, IdArray vids) const {
}
bool UnitGraph::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
const SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->HasEdgeBetween(etype, dst, src);
......@@ -854,7 +864,7 @@ bool UnitGraph::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) con
BoolArray UnitGraph::HasEdgesBetween(
dgl_type_t etype, IdArray src, IdArray dst) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
const SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->HasEdgesBetween(etype, dst, src);
......@@ -863,7 +873,7 @@ BoolArray UnitGraph::HasEdgesBetween(
}
IdArray UnitGraph::Predecessors(dgl_type_t etype, dgl_id_t dst) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
const SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->Successors(etype, dst);
......@@ -872,13 +882,13 @@ IdArray UnitGraph::Predecessors(dgl_type_t etype, dgl_id_t dst) const {
}
IdArray UnitGraph::Successors(dgl_type_t etype, dgl_id_t src) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
const SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
return ptr->Successors(etype, src);
}
IdArray UnitGraph::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
const SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->EdgeId(etype, dst, src);
......@@ -886,31 +896,41 @@ IdArray UnitGraph::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
return ptr->EdgeId(etype, src, dst);
}
EdgeArray UnitGraph::EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
EdgeArray UnitGraph::EdgeIdsAll(dgl_type_t etype, IdArray src, IdArray dst) const {
const SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC) {
EdgeArray edges = ptr->EdgeIds(etype, dst, src);
EdgeArray edges = ptr->EdgeIdsAll(etype, dst, src);
return EdgeArray{edges.dst, edges.src, edges.id};
} else {
return ptr->EdgeIds(etype, src, dst);
return ptr->EdgeIdsAll(etype, src, dst);
}
}
IdArray UnitGraph::EdgeIdsOne(dgl_type_t etype, IdArray src, IdArray dst) const {
const SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC) {
return ptr->EdgeIdsOne(etype, dst, src);
} else {
return ptr->EdgeIdsOne(etype, src, dst);
}
}
std::pair<dgl_id_t, dgl_id_t> UnitGraph::FindEdge(dgl_type_t etype, dgl_id_t eid) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCOO);
const SparseFormat fmt = SelectFormat(coo_code);
const auto ptr = GetFormat(fmt);
return ptr->FindEdge(etype, eid);
}
EdgeArray UnitGraph::FindEdges(dgl_type_t etype, IdArray eids) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCOO);
const SparseFormat fmt = SelectFormat(coo_code);
const auto ptr = GetFormat(fmt);
return ptr->FindEdges(etype, eids);
}
EdgeArray UnitGraph::InEdges(dgl_type_t etype, dgl_id_t vid) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
const SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC) {
const EdgeArray& ret = ptr->OutEdges(etype, vid);
......@@ -921,7 +941,7 @@ EdgeArray UnitGraph::InEdges(dgl_type_t etype, dgl_id_t vid) const {
}
EdgeArray UnitGraph::InEdges(dgl_type_t etype, IdArray vids) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
const SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC) {
const EdgeArray& ret = ptr->OutEdges(etype, vids);
......@@ -932,13 +952,13 @@ EdgeArray UnitGraph::InEdges(dgl_type_t etype, IdArray vids) const {
}
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, dgl_id_t vid) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
const SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
return ptr->OutEdges(etype, vid);
}
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, IdArray vids) const {
const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
const SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
return ptr->OutEdges(etype, vids);
}
......@@ -946,12 +966,12 @@ EdgeArray UnitGraph::OutEdges(dgl_type_t etype, IdArray vids) const {
EdgeArray UnitGraph::Edges(dgl_type_t etype, const std::string &order) const {
SparseFormat fmt;
if (order == std::string("eid")) {
fmt = SelectFormat(SparseFormat::kCOO);
fmt = SelectFormat(coo_code);
} else if (order.empty()) {
// arbitrary order
fmt = SelectFormat(SparseFormat::kAny);
fmt = SelectFormat(all_code);
} else if (order == std::string("srcdst")) {
fmt = SelectFormat(SparseFormat::kCSR);
fmt = SelectFormat(csr_code);
} else {
LOG(FATAL) << "Unsupported order request: " << order;
return {};
......@@ -965,7 +985,7 @@ EdgeArray UnitGraph::Edges(dgl_type_t etype, const std::string &order) const {
}
uint64_t UnitGraph::InDegree(dgl_type_t etype, dgl_id_t vid) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->OutDegree(etype, vid);
......@@ -974,7 +994,7 @@ uint64_t UnitGraph::InDegree(dgl_type_t etype, dgl_id_t vid) const {
}
DegreeArray UnitGraph::InDegrees(dgl_type_t etype, IdArray vids) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->OutDegrees(etype, vids);
......@@ -983,38 +1003,38 @@ DegreeArray UnitGraph::InDegrees(dgl_type_t etype, IdArray vids) const {
}
uint64_t UnitGraph::OutDegree(dgl_type_t etype, dgl_id_t vid) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
return ptr->OutDegree(etype, vid);
}
DegreeArray UnitGraph::OutDegrees(dgl_type_t etype, IdArray vids) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
return ptr->OutDegrees(etype, vids);
}
DGLIdIters UnitGraph::SuccVec(dgl_type_t etype, dgl_id_t vid) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
return ptr->SuccVec(etype, vid);
}
DGLIdIters32 UnitGraph::SuccVec32(dgl_type_t etype, dgl_id_t vid) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = std::dynamic_pointer_cast<CSR>(GetFormat(fmt));
CHECK_NOTNULL(ptr);
return ptr->SuccVec32(etype, vid);
}
DGLIdIters UnitGraph::OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
SparseFormat fmt = SelectFormat(csr_code);
const auto ptr = GetFormat(fmt);
return ptr->OutEdgeVec(etype, vid);
}
DGLIdIters UnitGraph::PredVec(dgl_type_t etype, dgl_id_t vid) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->SuccVec(etype, vid);
......@@ -1023,7 +1043,7 @@ DGLIdIters UnitGraph::PredVec(dgl_type_t etype, dgl_id_t vid) const {
}
DGLIdIters UnitGraph::InEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
SparseFormat fmt = SelectFormat(csc_code);
const auto ptr = GetFormat(fmt);
if (fmt == SparseFormat::kCSC)
return ptr->OutEdgeVec(etype, vid);
......@@ -1053,7 +1073,7 @@ std::vector<IdArray> UnitGraph::GetAdj(
HeteroSubgraph UnitGraph::VertexSubgraph(const std::vector<IdArray>& vids) const {
// We prefer to generate a subgraph from out-csr.
SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
SparseFormat fmt = SelectFormat(csr_code);
HeteroSubgraph sg = GetFormat(fmt)->VertexSubgraph(vids);
HeteroSubgraph ret;
......@@ -1083,7 +1103,7 @@ HeteroSubgraph UnitGraph::VertexSubgraph(const std::vector<IdArray>& vids) const
HeteroSubgraph UnitGraph::EdgeSubgraph(
const std::vector<IdArray>& eids, bool preserve_nodes) const {
SparseFormat fmt = SelectFormat(SparseFormat::kCOO);
SparseFormat fmt = SelectFormat(coo_code);
auto sg = GetFormat(fmt)->EdgeSubgraph(eids, preserve_nodes);
HeteroSubgraph ret;
......@@ -1114,7 +1134,7 @@ HeteroSubgraph UnitGraph::EdgeSubgraph(
HeteroGraphPtr UnitGraph::CreateFromCOO(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray row, IdArray col,
SparseFormat restrict_format) {
dgl_format_code_t formats) {
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(num_src, num_dst);
......@@ -1122,12 +1142,12 @@ HeteroGraphPtr UnitGraph::CreateFromCOO(
COOPtr coo(new COO(mg, num_src, num_dst, row, col));
return HeteroGraphPtr(
new UnitGraph(mg, nullptr, nullptr, coo, restrict_format));
new UnitGraph(mg, nullptr, nullptr, coo, formats));
}
HeteroGraphPtr UnitGraph::CreateFromCOO(
int64_t num_vtypes, const aten::COOMatrix& mat,
SparseFormat restrict_format) {
dgl_format_code_t formats) {
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(mat.num_rows, mat.num_cols);
......@@ -1135,51 +1155,51 @@ HeteroGraphPtr UnitGraph::CreateFromCOO(
COOPtr coo(new COO(mg, mat));
return HeteroGraphPtr(
new UnitGraph(mg, nullptr, nullptr, coo, restrict_format));
new UnitGraph(mg, nullptr, nullptr, coo, formats));
}
HeteroGraphPtr UnitGraph::CreateFromCSR(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids, SparseFormat restrict_format) {
IdArray indptr, IdArray indices, IdArray edge_ids, dgl_format_code_t formats) {
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(num_src, num_dst);
auto mg = CreateUnitGraphMetaGraph(num_vtypes);
CSRPtr csr(new CSR(mg, num_src, num_dst, indptr, indices, edge_ids));
return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, restrict_format));
return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, formats));
}
HeteroGraphPtr UnitGraph::CreateFromCSR(
int64_t num_vtypes, const aten::CSRMatrix& mat,
SparseFormat restrict_format) {
dgl_format_code_t formats) {
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(mat.num_rows, mat.num_cols);
auto mg = CreateUnitGraphMetaGraph(num_vtypes);
CSRPtr csr(new CSR(mg, mat));
return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, restrict_format));
return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, formats));
}
HeteroGraphPtr UnitGraph::CreateFromCSC(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids, SparseFormat restrict_format) {
IdArray indptr, IdArray indices, IdArray edge_ids, dgl_format_code_t formats) {
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(num_src, num_dst);
auto mg = CreateUnitGraphMetaGraph(num_vtypes);
CSRPtr csc(new CSR(mg, num_src, num_dst, indptr, indices, edge_ids));
return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, restrict_format));
return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, formats));
}
HeteroGraphPtr UnitGraph::CreateFromCSC(
int64_t num_vtypes, const aten::CSRMatrix& mat,
SparseFormat restrict_format) {
dgl_format_code_t formats) {
CHECK(num_vtypes == 1 || num_vtypes == 2);
if (num_vtypes == 1)
CHECK_EQ(mat.num_rows, mat.num_cols);
auto mg = CreateUnitGraphMetaGraph(num_vtypes);
CSRPtr csc(new CSR(mg, mat));
return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, restrict_format));
return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, formats));
}
HeteroGraphPtr UnitGraph::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
......@@ -1195,7 +1215,7 @@ HeteroGraphPtr UnitGraph::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
COOPtr new_coo =
(bg->coo_->defined())? COOPtr(new COO(bg->coo_->AsNumBits(bits))) : nullptr;
return HeteroGraphPtr(
new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->restrict_format_));
new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->formats_));
}
}
......@@ -1212,12 +1232,12 @@ HeteroGraphPtr UnitGraph::CopyTo(HeteroGraphPtr g, const DLContext& ctx) {
COOPtr new_coo =
(bg->coo_->defined())? COOPtr(new COO(bg->coo_->CopyTo(ctx))) : nullptr;
return HeteroGraphPtr(
new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->restrict_format_));
new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->formats_));
}
}
UnitGraph::UnitGraph(GraphPtr metagraph, CSRPtr in_csr, CSRPtr out_csr, COOPtr coo,
SparseFormat restrict_format)
dgl_format_code_t formats)
: BaseHeteroGraph(metagraph), in_csr_(in_csr), out_csr_(out_csr), coo_(coo) {
if (!in_csr_) {
in_csr_ = CSRPtr(new CSR());
......@@ -1228,31 +1248,11 @@ UnitGraph::UnitGraph(GraphPtr metagraph, CSRPtr in_csr, CSRPtr out_csr, COOPtr c
if (!coo_) {
coo_ = COOPtr(new COO());
}
restrict_format_ = AutoDetectFormat(in_csr_, out_csr_, coo_, restrict_format);
switch (restrict_format) {
case SparseFormat::kCSC:
in_csr_ = GetInCSR();
// cleaning other format
out_csr_ = out_csr_->defined() ? CSRPtr(new CSR()) : out_csr_;
coo_ = coo_->defined() ? COOPtr(new COO()) : coo_;
break;
case SparseFormat::kCSR:
out_csr_ = GetOutCSR();
// cleaning other format
in_csr_ = in_csr_->defined() ? CSRPtr(new CSR()) : in_csr_;
coo_ = coo_->defined() ? COOPtr(new COO()) : coo_;
break;
case SparseFormat::kCOO:
coo_ = GetCOO();
// cleaning other format
in_csr_ = in_csr_->defined() ? CSRPtr(new CSR()) : in_csr_;
out_csr_ = out_csr_->defined() ? CSRPtr(new CSR()) : out_csr_;
break;
default:
break;
}
formats_ = formats;
dgl_format_code_t created = GetCreatedFormats();
if ((formats | created) != formats)
LOG(FATAL) << "Graph created from formats: " << CodeToStr(created) <<
", which is not compatible with available formats: " << CodeToStr(formats);
CHECK(GetAny()) << "At least one graph structure should exist.";
}
......@@ -1263,7 +1263,7 @@ HeteroGraphPtr UnitGraph::CreateHomographFrom(
bool has_in_csr,
bool has_out_csr,
bool has_coo,
SparseFormat restrict_format) {
dgl_format_code_t formats) {
auto mg = CreateUnitGraphMetaGraph1();
CSRPtr in_csr_ptr = nullptr;
......@@ -1283,15 +1283,14 @@ HeteroGraphPtr UnitGraph::CreateHomographFrom(
else
coo_ptr = COOPtr(new COO());
return HeteroGraphPtr(new UnitGraph(mg, in_csr_ptr, out_csr_ptr, coo_ptr, restrict_format));
return HeteroGraphPtr(new UnitGraph(mg, in_csr_ptr, out_csr_ptr, coo_ptr, formats));
}
UnitGraph::CSRPtr UnitGraph::GetInCSR(bool inplace) const {
if (inplace)
if (restrict_format_ != SparseFormat::kAny &&
restrict_format_ != SparseFormat::kCSC)
LOG(FATAL) << "The graph have restricted sparse format " << GetRestrictFormat() <<
", cannot create CSC matrix.";
if (!(formats_ & csc_code))
LOG(FATAL) << "The graph have restricted sparse format " <<
CodeToStr(formats_) << ", cannot create CSC matrix.";
CSRPtr ret = in_csr_;
if (!in_csr_->defined()) {
if (out_csr_->defined()) {
......@@ -1318,10 +1317,9 @@ UnitGraph::CSRPtr UnitGraph::GetInCSR(bool inplace) const {
/* !\brief Return out csr. If not exist, transpose the other one.*/
UnitGraph::CSRPtr UnitGraph::GetOutCSR(bool inplace) const {
if (inplace)
if (restrict_format_ != SparseFormat::kAny &&
restrict_format_ != SparseFormat::kCSR)
LOG(FATAL) << "The graph have restricted sparse format " << GetRestrictFormat() <<
", cannot create CSR matrix.";
if (!(formats_ & csr_code))
LOG(FATAL) << "The graph have restricted sparse format " <<
CodeToStr(formats_) << ", cannot create CSR matrix.";
CSRPtr ret = out_csr_;
if (!out_csr_->defined()) {
if (in_csr_->defined()) {
......@@ -1347,10 +1345,9 @@ UnitGraph::CSRPtr UnitGraph::GetOutCSR(bool inplace) const {
/* !\brief Return coo. If not exist, create from csr.*/
UnitGraph::COOPtr UnitGraph::GetCOO(bool inplace) const {
if (inplace)
if (restrict_format_ != SparseFormat::kAny &&
restrict_format_ != SparseFormat::kCOO)
LOG(FATAL) << "The graph have restricted sparse format " << GetRestrictFormat() <<
", cannot create COO matrix.";
if (!(formats_ & coo_code))
LOG(FATAL) << "The graph have restricted sparse format " <<
CodeToStr(formats_) << ", cannot create COO matrix.";
COOPtr ret = coo_;
if (!coo_->defined()) {
if (in_csr_->defined()) {
......@@ -1395,78 +1392,65 @@ HeteroGraphPtr UnitGraph::GetAny() const {
}
}
dgl_format_code_t UnitGraph::GetFormatInUse() const {
dgl_format_code_t UnitGraph::GetCreatedFormats() const {
dgl_format_code_t ret = 0;
if (in_csr_->defined()) ret = ret | 1;
ret = ret << 1;
if (out_csr_->defined()) ret = ret | 1;
ret = ret << 1;
if (coo_->defined()) ret = ret | 1;
if (in_csr_->defined())
ret |= csc_code;
if (out_csr_->defined())
ret |= csr_code;
if (coo_->defined())
ret |= coo_code;
return ret;
}
dgl_format_code_t UnitGraph::GetAllowedFormats() const {
return formats_;
}
HeteroGraphPtr UnitGraph::GetFormat(SparseFormat format) const {
switch (format) {
case SparseFormat::kCSR:
return GetOutCSR();
case SparseFormat::kCSC:
return GetInCSR();
case SparseFormat::kCOO:
default:
return GetCOO();
case SparseFormat::kAny:
return GetAny();
default: // SparseFormat::kAuto
LOG(FATAL) << "Must specify a restrict format.";
return nullptr;
}
}
HeteroGraphPtr UnitGraph::GetGraphInFormat(SparseFormat restrict_format) const {
int64_t num_vtypes = NumVertexTypes();
switch (restrict_format) {
case SparseFormat::kCOO:
return CreateFromCOO(
num_vtypes, GetCOO(false)->adj(), restrict_format);
case SparseFormat::kCSC:
return CreateFromCSC(
num_vtypes, GetInCSR(false)->adj(), restrict_format);
case SparseFormat::kCSR:
return CreateFromCSR(
num_vtypes, GetOutCSR(false)->adj(), restrict_format);
case SparseFormat::kAny:
return HeteroGraphPtr(
HeteroGraphPtr UnitGraph::GetGraphInFormat(dgl_format_code_t formats) const {
if (formats == all_code)
return HeteroGraphPtr(
// TODO(xiangsx) Make it as graph storage.Clone()
new UnitGraph(meta_graph_,
(in_csr_->defined()) ? CSRPtr(new CSR(*in_csr_)) : nullptr,
(out_csr_->defined()) ? CSRPtr(new CSR(*out_csr_)) : nullptr,
(coo_->defined()) ? COOPtr(new COO(*coo_)) : nullptr,
restrict_format));
default: // SparseFormat::kAuto
LOG(FATAL) << "Must specify a restrict format.";
return nullptr;
}
}
SparseFormat UnitGraph::AutoDetectFormat(
CSRPtr in_csr, CSRPtr out_csr, COOPtr coo, SparseFormat restrict_format) const {
if (restrict_format != SparseFormat::kAuto)
return restrict_format;
if (coo && coo->defined() && coo->IsHypersparse())
return SparseFormat::kCOO;
return SparseFormat::kAny;
}
SparseFormat UnitGraph::SelectFormat(SparseFormat preferred_format) const {
if (restrict_format_ != SparseFormat::kAny)
return restrict_format_; // force to select the restricted format
else if (preferred_format != SparseFormat::kAny)
return preferred_format;
else if (in_csr_->defined())
return SparseFormat::kCSC;
else if (out_csr_->defined())
return SparseFormat::kCSR;
else
(in_csr_->defined())
? CSRPtr(new CSR(*in_csr_))
: nullptr,
(out_csr_->defined())
? CSRPtr(new CSR(*out_csr_))
: nullptr,
(coo_->defined())
? COOPtr(new COO(*coo_))
: nullptr,
formats));
int64_t num_vtypes = NumVertexTypes();
if (formats & coo_code)
return CreateFromCOO(num_vtypes, GetCOO(false)->adj(), formats);
if (formats & csr_code)
return CreateFromCSR(num_vtypes, GetOutCSR(false)->adj(), formats);
return CreateFromCSC(num_vtypes, GetInCSR(false)->adj(), formats);
}
SparseFormat UnitGraph::SelectFormat(dgl_format_code_t preferred_formats) const {
dgl_format_code_t common = preferred_formats & formats_;
dgl_format_code_t created = GetCreatedFormats();
if (common & created)
return DecodeFormat(common & created);
if (coo_->defined() && coo_->IsHypersparse()) // only allow coo for hypersparse graph.
return SparseFormat::kCOO;
if (common)
return DecodeFormat(common);
return DecodeFormat(created);
}
GraphPtr UnitGraph::AsImmutableGraph() const {
......@@ -1496,21 +1480,21 @@ GraphPtr UnitGraph::AsImmutableGraph() const {
HeteroGraphPtr UnitGraph::LineGraph(bool backtracking) const {
// TODO(xiangsx) currently we only support homogeneous graph
auto fmt = SelectFormat(SparseFormat::kAny);
auto fmt = SelectFormat(all_code);
switch (fmt) {
case SparseFormat::kCOO: {
return CreateFromCOO(1, aten::COOLineGraph(coo_->adj(), backtracking), SparseFormat::kAny);
return CreateFromCOO(1, aten::COOLineGraph(coo_->adj(), backtracking));
}
case SparseFormat::kCSR: {
const aten::CSRMatrix csr = GetCSRMatrix(0);
const aten::COOMatrix coo = aten::COOLineGraph(aten::CSRToCOO(csr, true), backtracking);
return CreateFromCOO(1, coo, SparseFormat::kAny);
return CreateFromCOO(1, coo);
}
case SparseFormat::kCSC: {
const aten::CSRMatrix csc = GetCSCMatrix(0);
const aten::CSRMatrix csr = aten::CSRTranspose(csc);
const aten::COOMatrix coo = aten::COOLineGraph(aten::CSRToCOO(csr, true), backtracking);
return CreateFromCOO(1, coo, SparseFormat::kAny);
return CreateFromCOO(1, coo);
}
default:
LOG(FATAL) << "None of CSC, CSR, COO exist";
......@@ -1526,11 +1510,32 @@ bool UnitGraph::Load(dmlc::Stream* fs) {
CHECK(fs->Read(&magicNum)) << "Invalid Magic Number";
CHECK_EQ(magicNum, kDGLSerialize_UnitGraphMagic) << "Invalid UnitGraph Data";
int64_t save_format_code, restrict_format_code;
int64_t save_format_code, formats_code;
CHECK(fs->Read(&save_format_code)) << "Invalid format";
CHECK(fs->Read(&restrict_format_code)) << "Invalid format";
restrict_format_ = static_cast<SparseFormat>(restrict_format_code);
CHECK(fs->Read(&formats_code)) << "Invalid format";
auto save_format = static_cast<SparseFormat>(save_format_code);
if (formats_code >> 32) {
formats_ = static_cast<dgl_format_code_t>(0xffffffff & formats_code);
} else {
// NOTE(zihao): to be compatible with old formats.
switch (formats_code & 0xffffffff) {
case 0:
formats_ = all_code;
break;
case 1:
formats_ = coo_code;
break;
case 2:
formats_ = csr_code;
break;
case 3:
formats_ = csc_code;
break;
default:
LOG(FATAL) << "Load graph failed, formats code " << formats_code <<
"not recognized.";
}
}
switch (save_format) {
case SparseFormat::kCOO:
......@@ -1567,9 +1572,9 @@ void UnitGraph::Save(dmlc::Stream* fs) const {
fs->Write(kDGLSerialize_UnitGraphMagic);
// Didn't write UnitGraph::meta_graph_, since it's included in the underlying
// sparse matrix
auto avail_fmt = SelectFormat(SparseFormat::kAny);
auto avail_fmt = SelectFormat(all_code);
fs->Write(static_cast<int64_t>(avail_fmt));
fs->Write(static_cast<int64_t>(restrict_format_));
fs->Write(static_cast<int64_t>(formats_ | 0x100000000));
switch (avail_fmt) {
case SparseFormat::kCOO:
fs->Write(GetCOO());
......@@ -1603,27 +1608,27 @@ UnitGraph::ToSimple() const {
IdArray count;
IdArray edge_map;
auto avail_fmt = SelectFormat(SparseFormat::kAny);
auto avail_fmt = SelectFormat(all_code);
switch (avail_fmt) {
case SparseFormat::kCOO: {
auto ret = aten::COOToSimple(coo_->adj());
auto ret = aten::COOToSimple(GetCOO()->adj());
count = std::get<1>(ret);
edge_map = std::get<2>(ret);
new_coo = COOPtr(new COO(coo_->meta_graph(), std::get<0>(ret)));
new_coo = COOPtr(new COO(meta_graph(), std::get<0>(ret)));
break;
}
case SparseFormat::kCSR: {
auto ret = aten::CSRToSimple(in_csr_->adj());
auto ret = aten::CSRToSimple(GetOutCSR()->adj());
count = std::get<1>(ret);
edge_map = std::get<2>(ret);
new_incsr = CSRPtr(new CSR(in_csr_->meta_graph(), std::get<0>(ret)));
new_outcsr = CSRPtr(new CSR(meta_graph(), std::get<0>(ret)));
break;
}
case SparseFormat::kCSC: {
auto ret = aten::CSRToSimple(out_csr_->adj());
auto ret = aten::CSRToSimple(GetInCSR()->adj());
count = std::get<1>(ret);
edge_map = std::get<2>(ret);
new_outcsr = CSRPtr(new CSR(out_csr_->meta_graph(), std::get<0>(ret)));
new_incsr = CSRPtr(new CSR(meta_graph(), std::get<0>(ret)));
break;
}
default:
......
src/graph/unit_graph.h
View file @ 44089c8b
......@@ -116,7 +116,9 @@ class UnitGraph : public BaseHeteroGraph {
IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override;
EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override;
EdgeArray EdgeIdsAll(dgl_type_t etype, IdArray src, IdArray dst) const override;
IdArray EdgeIdsOne(dgl_type_t etype, IdArray src, IdArray dst) const override;
std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override;
......@@ -172,31 +174,31 @@ class UnitGraph : public BaseHeteroGraph {
/*! \brief Create a graph from COO arrays */
static HeteroGraphPtr CreateFromCOO(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray row, IdArray col, SparseFormat restrict_format = SparseFormat::kAuto);
IdArray row, IdArray col, dgl_format_code_t formats = all_code);
static HeteroGraphPtr CreateFromCOO(
int64_t num_vtypes, const aten::COOMatrix& mat,
SparseFormat restrict_format = SparseFormat::kAuto);
dgl_format_code_t formats = all_code);
/*! \brief Create a graph from (out) CSR arrays */
static HeteroGraphPtr CreateFromCSR(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids,
SparseFormat restrict_format = SparseFormat::kAuto);
dgl_format_code_t formats = all_code);
static HeteroGraphPtr CreateFromCSR(
int64_t num_vtypes, const aten::CSRMatrix& mat,
SparseFormat restrict_format = SparseFormat::kAuto);
dgl_format_code_t formats = all_code);
/*! \brief Create a graph from (in) CSC arrays */
static HeteroGraphPtr CreateFromCSC(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids,
SparseFormat restrict_format = SparseFormat::kAuto);
dgl_format_code_t formats = all_code);
static HeteroGraphPtr CreateFromCSC(
int64_t num_vtypes, const aten::CSRMatrix& mat,
SparseFormat restrict_format = SparseFormat::kAuto);
dgl_format_code_t formats = all_code);
/*! \brief Convert the graph to use the given number of bits for storage */
static HeteroGraphPtr AsNumBits(HeteroGraphPtr g, uint8_t bits);
......@@ -237,16 +239,8 @@ class UnitGraph : public BaseHeteroGraph {
/*! \return Return the out-edge CSR in the matrix form */
aten::CSRMatrix GetCSRMatrix(dgl_type_t etype) const override;
/*! \brief some heuristic rules to determine the restrict format. */
SparseFormat AutoDetectFormat(
CSRPtr in_csr, CSRPtr out_csr, COOPtr coo, SparseFormat restrict_format) const;
SparseFormat SelectFormat(dgl_type_t etype, SparseFormat preferred_format) const override {
return SelectFormat(preferred_format);
}
std::string GetRestrictFormat() const override {
return ToStringSparseFormat(this->restrict_format_);
SparseFormat SelectFormat(dgl_type_t etype, dgl_format_code_t preferred_formats) const override {
return SelectFormat(preferred_formats);
}
/*!
......@@ -257,9 +251,11 @@ class UnitGraph : public BaseHeteroGraph {
*/
HeteroGraphPtr GetFormat(SparseFormat format) const;
dgl_format_code_t GetFormatInUse() const override;
dgl_format_code_t GetCreatedFormats() const override;
dgl_format_code_t GetAllowedFormats() const override;
HeteroGraphPtr GetGraphInFormat(SparseFormat restrict_format) const override;
HeteroGraphPtr GetGraphInFormat(dgl_format_code_t formats) const override;
/*! \return Load UnitGraph from stream, using CSRMatrix*/
bool Load(dmlc::Stream* fs);
......@@ -296,7 +292,7 @@ class UnitGraph : public BaseHeteroGraph {
* \param coo coo
*/
UnitGraph(GraphPtr metagraph, CSRPtr in_csr, CSRPtr out_csr, COOPtr coo,
SparseFormat restrict_format = SparseFormat::kAuto);
dgl_format_code_t formats = all_code);
/*!
* \brief constructor
......@@ -315,7 +311,7 @@ class UnitGraph : public BaseHeteroGraph {
bool has_in_csr,
bool has_out_csr,
bool has_coo,
SparseFormat restrict_format = SparseFormat::kAuto);
dgl_format_code_t formats = all_code);
/*! \return Return any existing format. */
HeteroGraphPtr GetAny() const;
......@@ -329,7 +325,7 @@ class UnitGraph : public BaseHeteroGraph {
* Otherwise, it will return whatever DGL thinks is the most appropriate given
* the arguments.
*/
SparseFormat SelectFormat(SparseFormat preferred_format) const;
SparseFormat SelectFormat(dgl_format_code_t preferred_formats) const;
/*! \return Whether the graph is hypersparse */
bool IsHypersparse() const;
......@@ -346,12 +342,8 @@ class UnitGraph : public BaseHeteroGraph {
COOPtr coo_;
/*!
* \brief Storage format restriction.
* If it is not ANY, then conversion is not allowed for graph queries.
*
* Note that GetInCSR/GetOutCSR/GetCOO() can still be called and the conversion will
* still be done if requested explicitly (e.g. in message passing).
*/
SparseFormat restrict_format_;
dgl_format_code_t formats_;
};
}; // namespace dgl
......
tests/backend/__init__.py
View file @ 44089c8b
......@@ -49,16 +49,10 @@ def randn(shape):
return copy_to(_randn(shape), _default_context)
def tensor(data, dtype=None):
if dtype is None:
if is_tensor(data):
data = zerocopy_to_numpy(data)
else:
data = np.array(data)
dtype = int64 if np.issubdtype(data.dtype, np.integer) else float32
return copy_to(_tensor(data, dtype), _default_context)
def arange(start, stop):
return copy_to(_arange(start, stop), _default_context)
def arange(start, stop, dtype=int64):
return copy_to(_arange(start, stop, dtype), _default_context)
def full(shape, fill_value, dtype, ctx=_default_context):
return _full(shape, fill_value, dtype, ctx)
......
tests/compute/test_basics.py
View file @ 44089c8b
......@@ -5,6 +5,7 @@ import scipy.sparse as ssp
import networkx as nx
from dgl import DGLGraph
from collections import defaultdict as ddict
import unittest
D = 5
reduce_msg_shapes = set()
......@@ -34,6 +35,7 @@ def generate_graph(grad=False):
g.add_edge(i, 9)
# add a back flow from 9 to 0
g.add_edge(9, 0)
g = g.to(F.ctx())
ncol = F.randn((10, D))
ecol = F.randn((17, D))
if grad:
......@@ -55,15 +57,15 @@ def test_batch_setter_getter():
assert F.allclose(g.ndata['h'], F.zeros((10, D)))
# pop nodes
old_len = len(g.ndata)
assert _pfc(g.pop_n_repr('h')) == [0.] * 10
g.ndata.pop('h')
assert len(g.ndata) == old_len - 1
g.ndata['h'] = F.zeros((10, D))
# set partial nodes
u = F.tensor([1, 3, 5])
u = F.tensor([1, 3, 5], g.idtype)
g.nodes[u].data['h'] = F.ones((3, D))
assert _pfc(g.ndata['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
# get partial nodes
u = F.tensor([1, 2, 3])
u = F.tensor([1, 2, 3], g.idtype)
assert _pfc(g.nodes[u].data['h']) == [1., 0., 1.]
'''
......@@ -91,39 +93,39 @@ def test_batch_setter_getter():
assert _pfc(g.edata['l']) == [0.] * 17
# pop edges
old_len = len(g.edata)
assert _pfc(g.pop_e_repr('l')) == [0.] * 17
g.edata.pop('l')
assert len(g.edata) == old_len - 1
g.edata['l'] = F.zeros((17, D))
# set partial edges (many-many)
u = F.tensor([0, 0, 2, 5, 9])
v = F.tensor([1, 3, 9, 9, 0])
u = F.tensor([0, 0, 2, 5, 9], g.idtype)
v = F.tensor([1, 3, 9, 9, 0], g.idtype)
g.edges[u, v].data['l'] = F.ones((5, D))
truth = [0.] * 17
truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
assert _pfc(g.edata['l']) == truth
# set partial edges (many-one)
u = F.tensor([3, 4, 6])
v = F.tensor([9])
u = F.tensor([3, 4, 6], g.idtype)
v = F.tensor([9], g.idtype)
g.edges[u, v].data['l'] = F.ones((3, D))
truth[5] = truth[7] = truth[11] = 1.
assert _pfc(g.edata['l']) == truth
# set partial edges (one-many)
u = F.tensor([0])
v = F.tensor([4, 5, 6])
u = F.tensor([0], g.idtype)
v = F.tensor([4, 5, 6], g.idtype)
g.edges[u, v].data['l'] = F.ones((3, D))
truth[6] = truth[8] = truth[10] = 1.
assert _pfc(g.edata['l']) == truth
# get partial edges (many-many)
u = F.tensor([0, 6, 0])
v = F.tensor([6, 9, 7])
u = F.tensor([0, 6, 0], g.idtype)
v = F.tensor([6, 9, 7], g.idtype)
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
# get partial edges (many-one)
u = F.tensor([5, 6, 7])
v = F.tensor([9])
u = F.tensor([5, 6, 7], g.idtype)
v = F.tensor([9], g.idtype)
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
# get partial edges (one-many)
u = F.tensor([0])
v = F.tensor([3, 4, 5])
u = F.tensor([0], g.idtype)
v = F.tensor([3, 4, 5], g.idtype)
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 1.]
......@@ -131,7 +133,7 @@ def test_batch_setter_autograd():
g = generate_graph(grad=True)
h1 = g.ndata['h']
# partial set
v = F.tensor([1, 2, 8])
v = F.tensor([1, 2, 8], g.idtype)
hh = F.attach_grad(F.zeros((len(v), D)))
with F.record_grad():
g.nodes[v].data['h'] = hh
......@@ -140,7 +142,7 @@ def test_batch_setter_autograd():
assert F.array_equal(F.grad(h1)[:,0], F.tensor([2., 0., 0., 2., 2., 2., 2., 2., 0., 2.]))
assert F.array_equal(F.grad(hh)[:,0], F.tensor([2., 2., 2.]))
def test_nx_conversion():
def _test_nx_conversion():
# check conversion between networkx and DGLGraph
def _check_nx_feature(nxg, nf, ef):
......@@ -191,7 +193,7 @@ def test_nx_conversion():
# convert to DGLGraph, nx graph has id in edge feature
# use id feature to test non-tensor copy
g.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'])
g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'])
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
......@@ -203,10 +205,10 @@ def test_nx_conversion():
assert F.allclose(g.ndata['n1'], n1)
# with id in nx edge feature, e1 should follow original order
assert F.allclose(g.edata['e1'], e1)
assert F.array_equal(g.get_e_repr()['id'], F.copy_to(F.arange(0, 4), F.cpu()))
assert F.array_equal(F.astype(g.edata['id'], F.int64), F.copy_to(F.arange(0, 4), F.cpu()))
# test conversion after modifying DGLGraph
g.pop_e_repr('id') # pop id so we don't need to provide id when adding edges
g.edata.pop('id') # pop id so we don't need to provide id when adding edges
new_n = F.randn((2, 3))
new_e = F.randn((3, 5))
g.add_nodes(2, data={'n1': new_n})
......@@ -225,8 +227,7 @@ def test_nx_conversion():
for _, _, attr in nxg.edges(data=True):
attr.pop('id')
# test with a new graph
g = DGLGraph()
g.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1'])
g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1'])
# check graph size
assert g.number_of_nodes() == 7
assert g.number_of_edges() == 7
......@@ -251,8 +252,7 @@ def test_nx_conversion():
for u, v, d in nxg.edges(data=True):
d['h'] = F.tensor([u, v])
g = dgl.DGLGraph()
g.from_networkx(nxg, node_attrs=['h'], edge_attrs=['h'])
g = dgl.from_networkx(nxg, node_attrs=['h'], edge_attrs=['h'])
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 4
assert g.has_edge_between(0, 1)
......@@ -261,48 +261,14 @@ def test_nx_conversion():
assert F.allclose(g.edata['h'], F.tensor([[1., 2.], [1., 2.],
[2., 3.], [2., 3.]]))
def test_batch_send():
g = generate_graph()
def _fmsg(edges):
assert tuple(F.shape(edges.src['h'])) == (5, D)
return {'m' : edges.src['h']}
g.register_message_func(_fmsg)
# many-many send
u = F.tensor([0, 0, 0, 0, 0])
v = F.tensor([1, 2, 3, 4, 5])
g.send((u, v))
# one-many send
u = F.tensor([0])
v = F.tensor([1, 2, 3, 4, 5])
g.send((u, v))
# many-one send
u = F.tensor([1, 2, 3, 4, 5])
v = F.tensor([9])
g.send((u, v))
def test_batch_recv():
# basic recv test
g = generate_graph()
g.register_message_func(message_func)
g.register_reduce_func(reduce_func)
g.register_apply_node_func(apply_node_func)
u = F.tensor([0, 0, 0, 4, 5, 6])
v = F.tensor([1, 2, 3, 9, 9, 9])
reduce_msg_shapes.clear()
g.send((u, v))
g.recv(F.unique(v))
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
def test_apply_nodes():
def _upd(nodes):
return {'h' : nodes.data['h'] * 2}
g = generate_graph()
g.register_apply_node_func(_upd)
old = g.ndata['h']
g.apply_nodes()
g.apply_nodes(_upd)
assert F.allclose(old * 2, g.ndata['h'])
u = F.tensor([0, 3, 4, 6])
u = F.tensor([0, 3, 4, 6], g.idtype)
g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
assert F.allclose(F.gather_row(g.ndata['h'], u), F.zeros((4, D)))
......@@ -310,27 +276,23 @@ def test_apply_edges():
def _upd(edges):
return {'w' : edges.data['w'] * 2}
g = generate_graph()
g.register_apply_edge_func(_upd)
old = g.edata['w']
g.apply_edges()
g.apply_edges(_upd)
assert F.allclose(old * 2, g.edata['w'])
u = F.tensor([0, 0, 0, 4, 5, 6])
v = F.tensor([1, 2, 3, 9, 9, 9])
u = F.tensor([0, 0, 0, 4, 5, 6], g.idtype)
v = F.tensor([1, 2, 3, 9, 9, 9], g.idtype)
g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
eid = F.tensor(g.edge_ids(u, v))
assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
def test_update_routines():
g = generate_graph()
g.register_message_func(message_func)
g.register_reduce_func(reduce_func)
g.register_apply_node_func(apply_node_func)
# send_and_recv
reduce_msg_shapes.clear()
u = [0, 0, 0, 4, 5, 6]
v = [1, 2, 3, 9, 9, 9]
g.send_and_recv((u, v))
g.send_and_recv((u, v), message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
try:
......@@ -340,107 +302,29 @@ def test_update_routines():
pass
# pull
v = F.tensor([1, 2, 3, 9])
v = F.tensor([1, 2, 3, 9], g.idtype)
reduce_msg_shapes.clear()
g.pull(v)
g.pull(v, message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
reduce_msg_shapes.clear()
# push
v = F.tensor([0, 1, 2, 3])
v = F.tensor([0, 1, 2, 3], g.idtype)
reduce_msg_shapes.clear()
g.push(v)
g.push(v, message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
reduce_msg_shapes.clear()
# update_all
reduce_msg_shapes.clear()
g.update_all()
g.update_all(message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
reduce_msg_shapes.clear()
def test_recv_0deg():
# test recv with 0deg nodes;
g = DGLGraph()
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype, ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(_init2, 'h')
# test#1: recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata.pop('h')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# test#2: recv only 0deg node is equal to apply
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv(0)
new = g.ndata.pop('h')
# 0deg check: equal to apply_nodes
assert F.allclose(new[0], 2 * old[0])
# non-0deg check: untouched
assert F.allclose(new[1], old[1])
def test_recv_0deg_newfld():
# test recv with 0deg nodes; the reducer also creates a new field
g = DGLGraph()
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h1' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h1' : nodes.data['h1'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype=dtype, ctx=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
# test#1: recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.set_n_initializer(_init2, 'h1')
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata.pop('h1')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# test#2: recv only 0deg node
old = F.randn((2, 5))
g.ndata['h'] = old
g.ndata['h1'] = F.full((2, 5), -1, F.int64) # this is necessary
g.send((0, 1))
g.recv(0)
new = g.ndata.pop('h1')
# 0deg check: fallback to apply
assert F.allclose(new[0], F.full_1d(5, -2, F.int64))
# non-0deg check: not changed
assert F.allclose(new[1], F.full_1d(5, -1, F.int64))
def test_update_all_0deg():
# test#1
g = DGLGraph()
g = g.to(F.ctx())
g.add_nodes(5)
g.add_edge(1, 0)
g.add_edge(2, 0)
......@@ -467,6 +351,7 @@ def test_update_all_0deg():
# test#2: graph with no edge
g = DGLGraph()
g = g.to(F.ctx())
g.add_nodes(5)
g.set_n_initializer(_init2, 'h')
g.ndata['h'] = old_repr
......@@ -477,6 +362,7 @@ def test_update_all_0deg():
def test_pull_0deg():
g = DGLGraph()
g = g.to(F.ctx())
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
......@@ -487,14 +373,11 @@ def test_pull_0deg():
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype, ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(_init2, 'h')
# test#1: pull both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.ndata['h'] = old
g.pull([0, 1])
g.pull([0, 1], _message, _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
......@@ -504,86 +387,19 @@ def test_pull_0deg():
# test#2: pull only 0deg node
old = F.randn((2, 5))
g.ndata['h'] = old
g.pull(0)
g.pull(0, _message, _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: fallback to apply
assert F.allclose(new[0], 2*old[0])
# non-0deg check: not touched
assert F.allclose(new[1], old[1])
def test_send_multigraph():
g = DGLGraph()
g.add_nodes(3)
g.add_edge(0, 1)
g.add_edge(0, 1)
g.add_edge(0, 1)
g.add_edge(2, 1)
def _message_a(edges):
return {'a': edges.data['a']}
def _message_b(edges):
return {'a': edges.data['a'] * 3}
def _reduce(nodes):
return {'a': F.max(nodes.mailbox['a'], 1)}
def answer(*args):
return F.max(F.stack(args, 0), 0)
# send by eid
old_repr = F.randn((4, 5))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2]))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2, 3], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
# send on multigraph
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(([0, 2], [1, 1]), _message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], F.max(old_repr, 0))
# consecutive send and send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send((2, 1), _message_a)
g.send([0, 1], message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0] * 3, old_repr[1] * 3, old_repr[3]))
# consecutive send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(0, message_func=_message_a)
g.send(1, message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[1] * 3))
# send_and_recv_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send_and_recv([0, 2, 3], message_func=_message_a, reduce_func=_reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
assert F.allclose(new_repr[[0, 2]], F.zeros((2, 5)))
def test_dynamic_addition():
N = 3
D = 1
g = DGLGraph()
g = g.to(F.ctx())
# Test node addition
g.add_nodes(N)
......@@ -613,15 +429,16 @@ def test_dynamic_addition():
def test_repr():
G = dgl.DGLGraph()
G.add_nodes(10)
G.add_edge(0, 1)
repr_string = G.__repr__()
g = dgl.DGLGraph()
g = g.to(F.ctx())
g.add_nodes(10)
g.add_edge(0, 1)
repr_string = g.__repr__()
print(repr_string)
G.ndata['x'] = F.zeros((10, 5))
G.add_edges([0, 1], 2)
G.edata['y'] = F.zeros((3, 4))
repr_string = G.__repr__()
g.ndata['x'] = F.zeros((10, 5))
g.add_edges([0, 1], 2)
g.edata['y'] = F.zeros((3, 4))
repr_string = g.__repr__()
print(repr_string)
......@@ -632,6 +449,7 @@ def test_group_apply_edges():
return {"norm_feat": normalized_feat}
g = DGLGraph()
g = g.to(F.ctx())
g.add_nodes(10)
g.add_edges(0, [1, 2, 3, 4, 5, 6, 7, 8])
g.add_edges(1, [2, 3, 4, 6, 7, 8])
......@@ -659,18 +477,21 @@ def test_group_apply_edges():
# GitHub issue #1036
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
def test_group_apply_edges2():
m = ssp.random(10, 10, 0.2)
g = DGLGraph(m, readonly=True)
g = g.to(F.ctx())
g.ndata['deg'] = g.in_degrees()
g.ndata['id'] = F.arange(0, g.number_of_nodes())
g.edata['id'] = F.arange(0, g.number_of_edges())
g.ndata['id'] = F.arange(0, g.number_of_nodes(), g.idtype)
g.edata['id'] = F.arange(0, g.number_of_edges(), g.idtype)
def apply(edges):
w = edges.data['id']
n_nodes, deg = w.shape
dst = edges.dst['id'][:, 0]
# TODO: tmp hack
eid1 = F.asnumpy(g.in_edges(dst, 'eid')).reshape(n_nodes, deg).sort(1)
eid2 = F.asnumpy(edges.data['id']).sort(1)
assert np.array_equal(eid1, eid2)
......@@ -681,6 +502,7 @@ def test_group_apply_edges2():
def test_local_var():
g = DGLGraph(nx.path_graph(5))
g = g.to(F.ctx())
g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
g.edata['w'] = F.zeros((g.number_of_edges(), 4))
# test override
......@@ -718,6 +540,7 @@ def test_local_var():
# test initializer1
g = DGLGraph()
g = g.to(F.ctx())
g.add_nodes(2)
g.add_edges([0, 1], [1, 1])
g.set_n_initializer(dgl.init.zero_initializer)
......@@ -740,6 +563,7 @@ def test_local_var():
def test_local_scope():
g = DGLGraph(nx.path_graph(5))
g = g.to(F.ctx())
g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
g.edata['w'] = F.zeros((g.number_of_edges(), 4))
# test override
......@@ -791,6 +615,7 @@ def test_local_scope():
# test initializer1
g = DGLGraph()
g = g.to(F.ctx())
g.add_nodes(2)
g.add_edges([0, 1], [1, 1])
g.set_n_initializer(dgl.init.zero_initializer)
......@@ -812,7 +637,7 @@ def test_local_scope():
foo(g)
if __name__ == '__main__':
test_nx_conversion()
#test_nx_conversion()
test_batch_setter_getter()
test_batch_setter_autograd()
test_batch_send()
......
tests/compute/test_batched_graph.py
View file @ 44089c8b
import dgl
import backend as F
import unittest
from test_utils import parametrize_dtype
def tree1():
def tree1(idtype):
"""Generate a tree
0
/ \
......@@ -11,7 +12,7 @@ def tree1():
3 4
Edges are from leaves to root.
"""
g = dgl.DGLGraph()
g = dgl.DGLGraph().astype(idtype).to(F.ctx())
g.add_nodes(5)
g.add_edge(3, 1)
g.add_edge(4, 1)
......@@ -21,7 +22,7 @@ def tree1():
g.edata['h'] = F.randn((4, 10))
return g
def tree2():
def tree2(idtype):
"""Generate a tree
1
/ \
......@@ -30,7 +31,7 @@ def tree2():
2 0
Edges are from leaves to root.
"""
g = dgl.DGLGraph()
g = dgl.DGLGraph().astype(idtype).to(F.ctx())
g.add_nodes(5)
g.add_edge(2, 4)
g.add_edge(0, 4)
......@@ -40,16 +41,17 @@ def tree2():
g.edata['h'] = F.randn((4, 10))
return g
def test_batch_unbatch():
t1 = tree1()
t2 = tree2()
@parametrize_dtype
def test_batch_unbatch(idtype):
t1 = tree1(idtype)
t2 = tree2(idtype)
bg = dgl.batch([t1, t2])
assert bg.number_of_nodes() == 10
assert bg.number_of_edges() == 8
assert bg.batch_size == 2
assert bg.batch_num_nodes == [5, 5]
assert bg.batch_num_edges == [4, 4]
assert F.allclose(bg.batch_num_nodes(), F.tensor([5, 5]))
assert F.allclose(bg.batch_num_edges(), F.tensor([4, 4]))
tt1, tt2 = dgl.unbatch(bg)
assert F.allclose(t1.ndata['h'], tt1.ndata['h'])
......@@ -57,16 +59,17 @@ def test_batch_unbatch():
assert F.allclose(t2.ndata['h'], tt2.ndata['h'])
assert F.allclose(t2.edata['h'], tt2.edata['h'])
def test_batch_unbatch1():
t1 = tree1()
t2 = tree2()
@parametrize_dtype
def test_batch_unbatch1(idtype):
t1 = tree1(idtype)
t2 = tree2(idtype)
b1 = dgl.batch([t1, t2])
b2 = dgl.batch([t2, b1])
assert b2.number_of_nodes() == 15
assert b2.number_of_edges() == 12
assert b2.batch_size == 3
assert b2.batch_num_nodes == [5, 5, 5]
assert b2.batch_num_edges == [4, 4, 4]
assert F.allclose(b2.batch_num_nodes(), F.tensor([5, 5, 5]))
assert F.allclose(b2.batch_num_edges(), F.tensor([4, 4, 4]))
s1, s2, s3 = dgl.unbatch(b2)
assert F.allclose(t2.ndata['h'], s1.ndata['h'])
......@@ -76,40 +79,14 @@ def test_batch_unbatch1():
assert F.allclose(t2.ndata['h'], s3.ndata['h'])
assert F.allclose(t2.edata['h'], s3.edata['h'])
# Test batching readonly graphs
t1.readonly()
t2.readonly()
t1_was_readonly = t1.is_readonly
t2_was_readonly = t2.is_readonly
bg = dgl.batch([t1, t2])
assert t1.is_readonly == t1_was_readonly
assert t2.is_readonly == t2_was_readonly
assert bg.number_of_nodes() == 10
assert bg.number_of_edges() == 8
assert bg.batch_size == 2
assert bg.batch_num_nodes == [5, 5]
assert bg.batch_num_edges == [4, 4]
rs1, rs2 = dgl.unbatch(bg)
assert F.allclose(rs1.edges()[0], t1.edges()[0])
assert F.allclose(rs1.edges()[1], t1.edges()[1])
assert F.allclose(rs2.edges()[0], t2.edges()[0])
assert F.allclose(rs2.edges()[1], t2.edges()[1])
assert F.allclose(rs1.nodes(), t1.nodes())
assert F.allclose(rs2.nodes(), t2.nodes())
assert F.allclose(t1.ndata['h'], rs1.ndata['h'])
assert F.allclose(t1.edata['h'], rs1.edata['h'])
assert F.allclose(t2.ndata['h'], rs2.ndata['h'])
assert F.allclose(t2.edata['h'], rs2.edata['h'])
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
def test_batch_unbatch_frame():
@parametrize_dtype
def test_batch_unbatch_frame(idtype):
"""Test module of node/edge frames of batched/unbatched DGLGraphs.
Also address the bug mentioned in https://github.com/dmlc/dgl/issues/1475.
"""
t1 = tree1()
t2 = tree2()
t1 = tree1(idtype)
t2 = tree2(idtype)
N1 = t1.number_of_nodes()
E1 = t1.number_of_edges()
N2 = t2.number_of_nodes()
......@@ -140,12 +117,13 @@ def test_batch_unbatch_frame():
assert F.allclose(_g2.ndata['h'], F.zeros((N2, D)))
assert F.allclose(_g2.edata['h'], F.zeros((E2, D)))
def test_batch_unbatch2():
@parametrize_dtype
def test_batch_unbatch2(idtype):
# test setting/getting features after batch
a = dgl.DGLGraph()
a = dgl.DGLGraph().astype(idtype).to(F.ctx())
a.add_nodes(4)
a.add_edges(0, [1, 2, 3])
b = dgl.DGLGraph()
b = dgl.DGLGraph().astype(idtype).to(F.ctx())
b.add_nodes(3)
b.add_edges(0, [1, 2])
c = dgl.batch([a, b])
......@@ -154,46 +132,31 @@ def test_batch_unbatch2():
assert F.allclose(c.ndata['h'], F.ones((7, 1)))
assert F.allclose(c.edata['w'], F.ones((5, 1)))
def test_batch_send_then_recv():
t1 = tree1()
t2 = tree2()
bg = dgl.batch([t1, t2])
bg.register_message_func(lambda edges: {'m' : edges.src['h']})
bg.register_reduce_func(lambda nodes: {'h' : F.sum(nodes.mailbox['m'], 1)})
u = [3, 4, 2 + 5, 0 + 5]
v = [1, 1, 4 + 5, 4 + 5]
bg.send((u, v))
bg.recv([1, 9]) # assuming recv takes in unique nodes
t1, t2 = dgl.unbatch(bg)
assert F.asnumpy(t1.ndata['h'][1]) == 7
assert F.asnumpy(t2.ndata['h'][4]) == 2
def test_batch_send_and_recv():
t1 = tree1()
t2 = tree2()
@parametrize_dtype
def test_batch_send_and_recv(idtype):
t1 = tree1(idtype)
t2 = tree2(idtype)
bg = dgl.batch([t1, t2])
bg.register_message_func(lambda edges: {'m' : edges.src['h']})
bg.register_reduce_func(lambda nodes: {'h' : F.sum(nodes.mailbox['m'], 1)})
_mfunc = lambda edges: {'m' : edges.src['h']}
_rfunc = lambda nodes: {'h' : F.sum(nodes.mailbox['m'], 1)}
u = [3, 4, 2 + 5, 0 + 5]
v = [1, 1, 4 + 5, 4 + 5]
bg.send_and_recv((u, v))
bg.send_and_recv((u, v), _mfunc, _rfunc)
t1, t2 = dgl.unbatch(bg)
assert F.asnumpy(t1.ndata['h'][1]) == 7
assert F.asnumpy(t2.ndata['h'][4]) == 2
def test_batch_propagate():
t1 = tree1()
t2 = tree2()
@parametrize_dtype
def test_batch_propagate(idtype):
t1 = tree1(idtype)
t2 = tree2(idtype)
bg = dgl.batch([t1, t2])
bg.register_message_func(lambda edges: {'m' : edges.src['h']})
bg.register_reduce_func(lambda nodes: {'h' : F.sum(nodes.mailbox['m'], 1)})
_mfunc = lambda edges: {'m' : edges.src['h']}
_rfunc = lambda nodes: {'h' : F.sum(nodes.mailbox['m'], 1)}
# get leaves.
order = []
......@@ -208,19 +171,20 @@ def test_batch_propagate():
v = [0, 0, 1 + 5, 1 + 5]
order.append((u, v))
bg.prop_edges(order)
bg.prop_edges(order, _mfunc, _rfunc)
t1, t2 = dgl.unbatch(bg)
assert F.asnumpy(t1.ndata['h'][0]) == 9
assert F.asnumpy(t2.ndata['h'][1]) == 5
def test_batched_edge_ordering():
g1 = dgl.DGLGraph()
@parametrize_dtype
def test_batched_edge_ordering(idtype):
g1 = dgl.DGLGraph().astype(idtype).to(F.ctx())
g1.add_nodes(6)
g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
e1 = F.randn((5, 10))
g1.edata['h'] = e1
g2 = dgl.DGLGraph()
g2 = dgl.DGLGraph().astype(idtype).to(F.ctx())
g2.add_nodes(6)
g2.add_edges([0, 1 ,2 ,5, 4 ,5], [1, 2, 3, 4, 3, 0])
e2 = F.randn((6, 10))
......@@ -230,14 +194,15 @@ def test_batched_edge_ordering():
r2 = g1.edata['h'][g1.edge_id(4, 5)]
assert F.array_equal(r1, r2)
def test_batch_no_edge():
g1 = dgl.DGLGraph()
@parametrize_dtype
def test_batch_no_edge(idtype):
g1 = dgl.DGLGraph().astype(idtype).to(F.ctx())
g1.add_nodes(6)
g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
g2 = dgl.DGLGraph()
g2 = dgl.DGLGraph().astype(idtype).to(F.ctx())
g2.add_nodes(6)
g2.add_edges([0, 1, 2, 5, 4, 5], [1 ,2 ,3, 4, 3, 0])
g3 = dgl.DGLGraph()
g3 = dgl.DGLGraph().astype(idtype).to(F.ctx())
g3.add_nodes(1) # no edges
g = dgl.batch([g1, g3, g2]) # should not throw an error
......
tests/compute/test_batched_heterograph.py
View file @ 44089c8b
......@@ -4,6 +4,7 @@ import unittest
from dgl.base import ALL
from utils import parametrize_dtype
from test_utils import check_graph_equal
def check_equivalence_between_heterographs(g1, g2, node_attrs=None, edge_attrs=None):
assert g1.ntypes == g2.ntypes
......@@ -19,8 +20,8 @@ def check_equivalence_between_heterographs(g1, g2, node_attrs=None, edge_attrs=N
for ety in g1.canonical_etypes:
assert g1.number_of_edges(ety) == g2.number_of_edges(ety)
src1, dst1, eid1 = g1.all_edges(etype=ety, form='all')
src2, dst2, eid2 = g2.all_edges(etype=ety, form='all')
src1, dst1, eid1 = g1.edges(etype=ety, form='all')
src2, dst2, eid2 = g2.edges(etype=ety, form='all')
assert F.allclose(src1, src2)
assert F.allclose(dst1, dst2)
assert F.allclose(eid1, eid2)
......@@ -40,20 +41,22 @@ def check_equivalence_between_heterographs(g1, g2, node_attrs=None, edge_attrs=N
assert F.allclose(g1.edges[ety].data[feat_name], g2.edges[ety].data[feat_name])
@parametrize_dtype
def test_batching_hetero_topology(index_dtype):
def test_topology(idtype):
"""Test batching two DGLHeteroGraphs where some nodes are isolated in some relations"""
g1 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'follows', 'developer'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0), (2, 1), (3, 1)]
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
g2 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'follows', 'developer'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0), (2, 1)]
}, index_dtype=index_dtype)
bg = dgl.batch_hetero([g1, g2])
}, idtype=idtype, device=F.ctx())
bg = dgl.batch([g1, g2])
assert bg.idtype == idtype
assert bg.device == F.ctx()
assert bg.ntypes == g2.ntypes
assert bg.etypes == g2.etypes
assert bg.canonical_etypes == g2.canonical_etypes
......@@ -61,72 +64,15 @@ def test_batching_hetero_topology(index_dtype):
# Test number of nodes
for ntype in bg.ntypes:
assert bg.batch_num_nodes(ntype) == [
print(ntype)
assert F.asnumpy(bg.batch_num_nodes(ntype)).tolist() == [
g1.number_of_nodes(ntype), g2.number_of_nodes(ntype)]
assert bg.number_of_nodes(ntype) == (
g1.number_of_nodes(ntype) + g2.number_of_nodes(ntype))
# Test number of edges
assert bg.batch_num_edges('plays') == [
g1.number_of_edges('plays'), g2.number_of_edges('plays')]
assert bg.number_of_edges('plays') == (
g1.number_of_edges('plays') + g2.number_of_edges('plays'))
for etype in bg.canonical_etypes:
assert bg.batch_num_edges(etype) == [
g1.number_of_edges(etype), g2.number_of_edges(etype)]
assert bg.number_of_edges(etype) == (
g1.number_of_edges(etype) + g2.number_of_edges(etype))
# Test relabeled nodes
for ntype in bg.ntypes:
assert list(F.asnumpy(bg.nodes(ntype))) == list(range(bg.number_of_nodes(ntype)))
# Test relabeled edges
src, dst = bg.all_edges(etype=('user', 'follows', 'user'))
assert list(F.asnumpy(src)) == [0, 1, 4, 5]
assert list(F.asnumpy(dst)) == [1, 2, 5, 6]
src, dst = bg.all_edges(etype=('user', 'follows', 'developer'))
assert list(F.asnumpy(src)) == [0, 1, 4, 5]
assert list(F.asnumpy(dst)) == [1, 2, 4, 5]
src, dst, eid = bg.all_edges(etype='plays', form='all')
assert list(F.asnumpy(src)) == [0, 1, 2, 3, 4, 5, 6]
assert list(F.asnumpy(dst)) == [0, 0, 1, 1, 2, 2, 3]
assert list(F.asnumpy(eid)) == [0, 1, 2, 3, 4, 5, 6]
# Test unbatching graphs
g3, g4 = dgl.unbatch_hetero(bg)
check_equivalence_between_heterographs(g1, g3)
check_equivalence_between_heterographs(g2, g4)
"""Test batching two DGLHeteroGraphs with csr format"""
g1 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'follows', 'developer'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0), (2, 1), (3, 1)]
}, index_dtype=index_dtype, restrict_format='csr')
g2 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'follows', 'developer'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0), (2, 1)]
}, index_dtype=index_dtype, restrict_format='csr')
bg = dgl.batch_hetero([g1, g2])
# Test number of nodes
for ntype in bg.ntypes:
assert bg.batch_num_nodes(ntype) == [
g1.number_of_nodes(ntype), g2.number_of_nodes(ntype)]
assert bg.number_of_nodes(ntype) == (
g1.number_of_nodes(ntype) + g2.number_of_nodes(ntype))
# Test number of edges
assert bg.batch_num_edges('plays') == [
g1.number_of_edges('plays'), g2.number_of_edges('plays')]
assert bg.number_of_edges('plays') == (
g1.number_of_edges('plays') + g2.number_of_edges('plays'))
for etype in bg.canonical_etypes:
assert bg.batch_num_edges(etype) == [
assert F.asnumpy(bg.batch_num_edges(etype)).tolist() == [
g1.number_of_edges(etype), g2.number_of_edges(etype)]
assert bg.number_of_edges(etype) == (
g1.number_of_edges(etype) + g2.number_of_edges(etype))
......@@ -136,92 +82,52 @@ def test_batching_hetero_topology(index_dtype):
assert list(F.asnumpy(bg.nodes(ntype))) == list(range(bg.number_of_nodes(ntype)))
# Test relabeled edges
src, dst = bg.all_edges(etype=('user', 'follows', 'user'))
src, dst = bg.edges(etype=('user', 'follows', 'user'))
assert list(F.asnumpy(src)) == [0, 1, 4, 5]
assert list(F.asnumpy(dst)) == [1, 2, 5, 6]
src, dst = bg.all_edges(etype=('user', 'follows', 'developer'))
src, dst = bg.edges(etype=('user', 'follows', 'developer'))
assert list(F.asnumpy(src)) == [0, 1, 4, 5]
assert list(F.asnumpy(dst)) == [1, 2, 4, 5]
src, dst, eid = bg.all_edges(etype='plays', form='all')
src, dst, eid = bg.edges(etype='plays', form='all')
assert list(F.asnumpy(src)) == [0, 1, 2, 3, 4, 5, 6]
assert list(F.asnumpy(dst)) == [0, 0, 1, 1, 2, 2, 3]
assert list(F.asnumpy(eid)) == [0, 1, 2, 3, 4, 5, 6]
# Test unbatching graphs
g3, g4 = dgl.unbatch_hetero(bg)
g3, g4 = dgl.unbatch(bg)
check_equivalence_between_heterographs(g1, g3)
check_equivalence_between_heterographs(g2, g4)
"""Test batching two DGLHeteroGraphs with csc"""
g1 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'follows', 'developer'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0), (2, 1), (3, 1)]
}, index_dtype=index_dtype, restrict_format='csc')
g2 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'follows', 'developer'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0), (2, 1)]
}, index_dtype=index_dtype, restrict_format='csc')
bg = dgl.batch_hetero([g1, g2])
# Test number of nodes
for ntype in bg.ntypes:
assert bg.batch_num_nodes(ntype) == [
g1.number_of_nodes(ntype), g2.number_of_nodes(ntype)]
assert bg.number_of_nodes(ntype) == (
g1.number_of_nodes(ntype) + g2.number_of_nodes(ntype))
# Test number of edges
assert bg.batch_num_edges('plays') == [
g1.number_of_edges('plays'), g2.number_of_edges('plays')]
assert bg.number_of_edges('plays') == (
g1.number_of_edges('plays') + g2.number_of_edges('plays'))
for etype in bg.canonical_etypes:
assert bg.batch_num_edges(etype) == [
g1.number_of_edges(etype), g2.number_of_edges(etype)]
assert bg.number_of_edges(etype) == (
g1.number_of_edges(etype) + g2.number_of_edges(etype))
# Test relabeled nodes
for ntype in bg.ntypes:
assert list(F.asnumpy(bg.nodes(ntype))) == list(range(bg.number_of_nodes(ntype)))
# Test relabeled edges
src, dst = bg.all_edges(etype=('user', 'follows', 'user'))
assert list(F.asnumpy(src)) == [0, 1, 4, 5]
assert list(F.asnumpy(dst)) == [1, 2, 5, 6]
src, dst = bg.all_edges(etype=('user', 'follows', 'developer'))
assert list(F.asnumpy(src)) == [0, 1, 4, 5]
assert list(F.asnumpy(dst)) == [1, 2, 4, 5]
src, dst, eid = bg.all_edges(etype='plays', form='all')
assert list(F.asnumpy(src)) == [0, 1, 2, 3, 4, 5, 6]
assert list(F.asnumpy(dst)) == [0, 0, 1, 1, 2, 2, 3]
assert list(F.asnumpy(eid)) == [0, 1, 2, 3, 4, 5, 6]
# Test dtype cast
if idtype == "int32":
bg_cast = bg.long()
else:
bg_cast = bg.int()
assert bg.batch_size == bg_cast.batch_size
# Test unbatching graphs
g3, g4 = dgl.unbatch_hetero(bg)
check_equivalence_between_heterographs(g1, g3)
check_equivalence_between_heterographs(g2, g4)
# Test local var
bg_local = bg.local_var()
assert bg.batch_size == bg_local.batch_size
@parametrize_dtype
def test_batching_hetero_and_batched_hetero_topology(index_dtype):
def test_batching_batched(idtype):
"""Test batching a DGLHeteroGraph and a BatchedDGLHeteroGraph."""
g1 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0)]
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
g2 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0)]
}, index_dtype=index_dtype)
bg1 = dgl.batch_hetero([g1, g2])
}, idtype=idtype, device=F.ctx())
bg1 = dgl.batch([g1, g2])
g3 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1)],
('user', 'plays', 'game'): [(1, 0)]
}, index_dtype=index_dtype)
bg2 = dgl.batch_hetero([bg1, g3])
}, idtype=idtype, device=F.ctx())
bg2 = dgl.batch([bg1, g3])
assert bg2.idtype == idtype
assert bg2.device == F.ctx()
assert bg2.ntypes == g3.ntypes
assert bg2.etypes == g3.etypes
assert bg2.canonical_etypes == g3.canonical_etypes
......@@ -229,20 +135,14 @@ def test_batching_hetero_and_batched_hetero_topology(index_dtype):
# Test number of nodes
for ntype in bg2.ntypes:
assert bg2.batch_num_nodes(ntype) == [
assert F.asnumpy(bg2.batch_num_nodes(ntype)).tolist() == [
g1.number_of_nodes(ntype), g2.number_of_nodes(ntype), g3.number_of_nodes(ntype)]
assert bg2.number_of_nodes(ntype) == (
g1.number_of_nodes(ntype) + g2.number_of_nodes(ntype) + g3.number_of_nodes(ntype))
# Test number of edges
for etype in bg2.etypes:
assert bg2.batch_num_edges(etype) == [
g1.number_of_edges(etype), g2.number_of_edges(etype), g3.number_of_edges(etype)]
assert bg2.number_of_edges(etype) == (
g1.number_of_edges(etype) + g2.number_of_edges(etype) + g3.number_of_edges(etype))
for etype in bg2.canonical_etypes:
assert bg2.batch_num_edges(etype) == [
assert F.asnumpy(bg2.batch_num_edges(etype)).tolist() == [
g1.number_of_edges(etype), g2.number_of_edges(etype), g3.number_of_edges(etype)]
assert bg2.number_of_edges(etype) == (
g1.number_of_edges(etype) + g2.number_of_edges(etype) + g3.number_of_edges(etype))
......@@ -252,26 +152,26 @@ def test_batching_hetero_and_batched_hetero_topology(index_dtype):
assert list(F.asnumpy(bg2.nodes(ntype))) == list(range(bg2.number_of_nodes(ntype)))
# Test relabeled edges
src, dst = bg2.all_edges(etype='follows')
src, dst = bg2.edges(etype='follows')
assert list(F.asnumpy(src)) == [0, 1, 3, 4, 6]
assert list(F.asnumpy(dst)) == [1, 2, 4, 5, 7]
src, dst = bg2.all_edges(etype='plays')
src, dst = bg2.edges(etype='plays')
assert list(F.asnumpy(src)) == [0, 1, 3, 4, 7]
assert list(F.asnumpy(dst)) == [0, 0, 1, 1, 2]
# Test unbatching graphs
g4, g5, g6 = dgl.unbatch_hetero(bg2)
g4, g5, g6 = dgl.unbatch(bg2)
check_equivalence_between_heterographs(g1, g4)
check_equivalence_between_heterographs(g2, g5)
check_equivalence_between_heterographs(g3, g6)
@parametrize_dtype
def test_batched_features(index_dtype):
def test_features(idtype):
"""Test the features of batched DGLHeteroGraphs"""
g1 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0)]
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
g1.nodes['user'].data['h1'] = F.tensor([[0.], [1.], [2.]])
g1.nodes['user'].data['h2'] = F.tensor([[3.], [4.], [5.]])
g1.nodes['game'].data['h1'] = F.tensor([[0.]])
......@@ -283,7 +183,7 @@ def test_batched_features(index_dtype):
g2 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0)]
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
g2.nodes['user'].data['h1'] = F.tensor([[0.], [1.], [2.]])
g2.nodes['user'].data['h2'] = F.tensor([[3.], [4.], [5.]])
g2.nodes['game'].data['h1'] = F.tensor([[0.]])
......@@ -292,13 +192,8 @@ def test_batched_features(index_dtype):
g2.edges['follows'].data['h2'] = F.tensor([[2.], [3.]])
g2.edges['plays'].data['h1'] = F.tensor([[0.], [1.]])
bg = dgl.batch_hetero([g1, g2],
node_attrs=ALL,
edge_attrs={
('user', 'follows', 'user'): 'h1',
('user', 'plays', 'game'): None
})
# test default setting
bg = dgl.batch([g1, g2])
assert F.allclose(bg.nodes['user'].data['h1'],
F.cat([g1.nodes['user'].data['h1'], g2.nodes['user'].data['h1']], dim=0))
assert F.allclose(bg.nodes['user'].data['h2'],
......@@ -309,27 +204,48 @@ def test_batched_features(index_dtype):
F.cat([g1.nodes['game'].data['h2'], g2.nodes['game'].data['h2']], dim=0))
assert F.allclose(bg.edges['follows'].data['h1'],
F.cat([g1.edges['follows'].data['h1'], g2.edges['follows'].data['h1']], dim=0))
assert 'h2' not in bg.edges['follows'].data.keys()
assert 'h1' not in bg.edges['plays'].data.keys()
assert F.allclose(bg.edges['follows'].data['h2'],
F.cat([g1.edges['follows'].data['h2'], g2.edges['follows'].data['h2']], dim=0))
assert F.allclose(bg.edges['plays'].data['h1'],
F.cat([g1.edges['plays'].data['h1'], g2.edges['plays'].data['h1']], dim=0))
# test specifying ndata/edata
bg = dgl.batch([g1, g2], ndata=['h2'], edata=['h1'])
assert F.allclose(bg.nodes['user'].data['h2'],
F.cat([g1.nodes['user'].data['h2'], g2.nodes['user'].data['h2']], dim=0))
assert F.allclose(bg.nodes['game'].data['h2'],
F.cat([g1.nodes['game'].data['h2'], g2.nodes['game'].data['h2']], dim=0))
assert F.allclose(bg.edges['follows'].data['h1'],
F.cat([g1.edges['follows'].data['h1'], g2.edges['follows'].data['h1']], dim=0))
assert F.allclose(bg.edges['plays'].data['h1'],
F.cat([g1.edges['plays'].data['h1'], g2.edges['plays'].data['h1']], dim=0))
assert 'h1' not in bg.nodes['user'].data
assert 'h1' not in bg.nodes['game'].data
assert 'h2' not in bg.edges['follows'].data
# Test unbatching graphs
g3, g4 = dgl.unbatch_hetero(bg)
g3, g4 = dgl.unbatch(bg)
check_equivalence_between_heterographs(
g1, g3,
node_attrs={'user': ['h1', 'h2'], 'game': ['h1', 'h2']},
node_attrs={'user': ['h2'], 'game': ['h2']},
edge_attrs={('user', 'follows', 'user'): ['h1']})
check_equivalence_between_heterographs(
g2, g4,
node_attrs={'user': ['h1', 'h2'], 'game': ['h1', 'h2']},
node_attrs={'user': ['h2'], 'game': ['h2']},
edge_attrs={('user', 'follows', 'user'): ['h1']})
# test legacy
bg = dgl.batch([g1, g2], edge_attrs=['h1'])
assert 'h2' not in bg.edges['follows'].data.keys()
@unittest.skipIf(F.backend_name == 'mxnet', reason="MXNet does not support split array with zero-length segment.")
@parametrize_dtype
def test_batching_with_zero_nodes_edges(index_dtype):
def test_empty_relation(idtype):
"""Test the features of batched DGLHeteroGraphs"""
g1 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): []
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
g1.nodes['user'].data['h1'] = F.tensor([[0.], [1.], [2.]])
g1.nodes['user'].data['h2'] = F.tensor([[3.], [4.], [5.]])
g1.edges['follows'].data['h1'] = F.tensor([[0.], [1.]])
......@@ -338,7 +254,7 @@ def test_batching_with_zero_nodes_edges(index_dtype):
g2 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0)]
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
g2.nodes['user'].data['h1'] = F.tensor([[0.], [1.], [2.]])
g2.nodes['user'].data['h2'] = F.tensor([[3.], [4.], [5.]])
g2.nodes['game'].data['h1'] = F.tensor([[0.]])
......@@ -347,8 +263,19 @@ def test_batching_with_zero_nodes_edges(index_dtype):
g2.edges['follows'].data['h2'] = F.tensor([[2.], [3.]])
g2.edges['plays'].data['h1'] = F.tensor([[0.], [1.]])
bg = dgl.batch_hetero([g1, g2])
bg = dgl.batch([g1, g2])
# Test number of nodes
for ntype in bg.ntypes:
assert F.asnumpy(bg.batch_num_nodes(ntype)).tolist() == [
g1.number_of_nodes(ntype), g2.number_of_nodes(ntype)]
# Test number of edges
for etype in bg.canonical_etypes:
assert F.asnumpy(bg.batch_num_edges(etype)).tolist() == [
g1.number_of_edges(etype), g2.number_of_edges(etype)]
# Test features
assert F.allclose(bg.nodes['user'].data['h1'],
F.cat([g1.nodes['user'].data['h1'], g2.nodes['user'].data['h1']], dim=0))
assert F.allclose(bg.nodes['user'].data['h2'],
......@@ -360,7 +287,7 @@ def test_batching_with_zero_nodes_edges(index_dtype):
assert F.allclose(bg.edges['plays'].data['h1'], g2.edges['plays'].data['h1'])
# Test unbatching graphs
g3, g4 = dgl.unbatch_hetero(bg)
g3, g4 = dgl.unbatch(bg)
check_equivalence_between_heterographs(
g1, g3,
node_attrs={'user': ['h1', 'h2'], 'game': ['h1', 'h2']},
......@@ -373,51 +300,36 @@ def test_batching_with_zero_nodes_edges(index_dtype):
# Test graphs without edges
g1 = dgl.bipartite([], 'u', 'r', 'v', num_nodes=(0, 4))
g2 = dgl.bipartite([], 'u', 'r', 'v', num_nodes=(1, 5))
g2.nodes['u'].data['x'] = F.tensor([1])
dgl.batch_hetero([g1, g2])
dgl.batch([g1, g2])
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
@parametrize_dtype
def test_to_device(index_dtype):
g1 = dgl.heterograph({
('user', 'plays', 'game'): [(0, 0), (1, 1)]
}, index_dtype=index_dtype)
g1.nodes['user'].data['h1'] = F.copy_to(F.tensor([[0.], [1.]]), F.cpu())
g1.nodes['user'].data['h2'] = F.copy_to(F.tensor([[3.], [4.]]), F.cpu())
g1.edges['plays'].data['h1'] = F.copy_to(F.tensor([[2.], [3.]]), F.cpu())
g2 = dgl.heterograph({
('user', 'plays', 'game'): [(0, 0), (1, 0)]
}, index_dtype=index_dtype)
g2.nodes['user'].data['h1'] = F.copy_to(F.tensor([[1.], [2.]]), F.cpu())
g2.nodes['user'].data['h2'] = F.copy_to(F.tensor([[4.], [5.]]), F.cpu())
g2.edges['plays'].data['h1'] = F.copy_to(F.tensor([[0.], [1.]]), F.cpu())
bg = dgl.batch_hetero([g1, g2])
if F.is_cuda_available():
bg1 = bg.to(F.cuda())
assert bg1 is not None
assert bg.batch_size == bg1.batch_size
assert bg.batch_num_nodes('user') == bg1.batch_num_nodes('user')
assert bg.batch_num_edges('plays') == bg1.batch_num_edges('plays')
# set feature
g1 = dgl.heterograph({
('user', 'plays', 'game'): [(0, 0), (1, 1)]
}, index_dtype=index_dtype)
g2 = dgl.heterograph({
('user', 'plays', 'game'): [(0, 0), (1, 0)]
}, index_dtype=index_dtype)
bg = dgl.batch_hetero([g1, g2])
if F.is_cuda_available():
bg1 = bg.to(F.cuda())
bg1.nodes['user'].data['test'] = F.copy_to(F.tensor([0,1,2,3]), F.cuda())
bg1.edata['test'] = F.copy_to(F.tensor([0,1,2,3]), F.cuda())
def test_unbatch2(idtype):
# batch 3 graphs but unbatch to 2
g1 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
g2 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
g3 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
bg = dgl.batch([g1, g2, g3])
bnn = F.tensor([8, 4])
bne = F.tensor([6, 3])
f1, f2 = dgl.unbatch(bg, node_split=bnn, edge_split=bne)
u, v = f1.edges(order='eid')
assert F.allclose(u, F.tensor([0, 1, 2, 4, 5, 6]))
assert F.allclose(v, F.tensor([1, 2, 3, 5, 6, 7]))
u, v = f2.edges(order='eid')
assert F.allclose(u, F.tensor([0, 1, 2]))
assert F.allclose(v, F.tensor([1, 2, 3]))
# batch 2 but unbatch to 3
bg = dgl.batch([f1, f2])
gg1, gg2, gg3 = dgl.unbatch(bg, F.tensor([4, 4, 4]), F.tensor([3, 3, 3]))
check_graph_equal(g1, gg1)
check_graph_equal(g2, gg2)
check_graph_equal(g3, gg3)
if __name__ == '__main__':
test_batching_hetero_topology('int32')
test_batching_hetero_and_batched_hetero_topology('int32')
test_batched_features('int32')
test_batching_with_zero_nodes_edges('int32')
test_to_device('int32')
#test_topology('int32')
#test_batching_batched('int32')
#test_batched_features('int32')
#test_empty_relation('int32')
#test_to_device('int32')
pass
tests/compute/test_function.py deleted 100644 → 0
View file @ 015acfd2
import dgl
import dgl.function as fn
import backend as F
def generate_graph():
g = dgl.DGLGraph()
g.add_nodes(10) # 10 nodes.
h = F.astype(F.arange(1, 11), F.float32)
g.ndata['h'] = h
# create a graph where 0 is the source and 9 is the sink
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
# add a back flow from 9 to 0
g.add_edge(9, 0)
h = F.tensor([1., 2., 1., 3., 1., 4., 1., 5., 1., 6.,\
1., 7., 1., 8., 1., 9., 10.])
g.edata['h'] = h
return g
def reducer_both(nodes):
return {'out' : F.sum(nodes.mailbox['m'], 1)}
def test_copy_src():
# copy_src with both fields
g = generate_graph()
g.register_message_func(fn.copy_src(src='h', out='m'))
g.register_reduce_func(reducer_both)
# test with update_all
g.update_all()
assert F.allclose(g.ndata.pop('out'),
F.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# test with send and then recv
g.send()
g.recv()
assert F.allclose(g.ndata.pop('out'),
F.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
def test_copy_edge():
# copy_edge with both fields
g = generate_graph()
g.register_message_func(fn.copy_edge(edge='h', out='m'))
g.register_reduce_func(reducer_both)
# test with update_all
g.update_all()
assert F.allclose(g.ndata.pop('out'),
F.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# test with send and then recv
g.send()
g.recv()
assert F.allclose(g.ndata.pop('out'),
F.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
def test_src_mul_edge():
# src_mul_edge with all fields
g = generate_graph()
g.register_message_func(fn.src_mul_edge(src='h', edge='h', out='m'))
g.register_reduce_func(reducer_both)
# test with update_all
g.update_all()
assert F.allclose(g.ndata.pop('out'),
F.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
# test with send and then recv
g.send()
g.recv()
assert F.allclose(g.ndata.pop('out'),
F.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
if __name__ == '__main__':
test_copy_src()
test_copy_edge()
test_src_mul_edge()
tests/compute/test_graph.py
View file @ 44089c8b
import math
import numbers
import numpy as np
import scipy.sparse as sp
import networkx as nx
......@@ -207,19 +208,16 @@ def test_query():
eids = g.edge_ids([4,0], [4,9])
assert eids.shape[0] == 2
eid = g.edge_id(4, 4)
assert isinstance(eid, int)
with pytest.raises(AssertionError):
assert isinstance(eid, numbers.Number)
with pytest.raises(DGLError):
eids = g.edge_ids([9,0], [4,9])
with pytest.raises(AssertionError):
with pytest.raises(DGLError):
eid = g.edge_id(4, 5)
g.add_edge(0, 4)
with pytest.raises(AssertionError):
eids = g.edge_ids([0,0], [4,9])
with pytest.raises(AssertionError):
eid = g.edge_id(0, 4)
eids = g.edge_ids([0,0], [4,9])
eid = g.edge_id(0, 4)
_test(gen_by_mutation())
_test(gen_from_data(elist_input(), False, False))
......@@ -236,6 +234,7 @@ def test_query():
def test_mutation():
g = dgl.DGLGraph()
g = g.to(F.ctx())
# test add nodes with data
g.add_nodes(5)
g.add_nodes(5, {'h' : F.ones((5, 2))})
......@@ -248,15 +247,6 @@ def test_mutation():
g.add_edges([0, 1], [1, 2], {'m' : F.ones((2, 2))})
ans = F.cat([F.zeros((2, 2)), F.ones((2, 2))], 0)
assert F.allclose(ans, g.edata['m'])
# test clear and add again
g.clear()
g.add_nodes(5)
g.ndata['h'] = 3 * F.ones((5, 2))
assert F.allclose(3 * F.ones((5, 2)), g.ndata['h'])
g.init_ndata('h1', (g.number_of_nodes(), 3), 'float32')
assert F.allclose(F.zeros((g.number_of_nodes(), 3)), g.ndata['h1'])
g.init_edata('h2', (g.number_of_edges(), 3), 'float32')
assert F.allclose(F.zeros((g.number_of_edges(), 3)), g.edata['h2'])
def test_scipy_adjmat():
g = dgl.DGLGraph()
......@@ -315,63 +305,6 @@ def test_incmat():
[0., 1., 0., -1., 0.],
[0., 0., 1., 1., 0.]]))
def test_readonly():
g = dgl.DGLGraph()
g.add_nodes(5)
g.add_edges([0, 1, 2, 3], [1, 2, 3, 4])
g.ndata['x'] = F.zeros((5, 3))
g.edata['x'] = F.zeros((4, 4))
g.readonly(False)
assert g._graph.is_readonly() == False
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
g.readonly()
assert g._graph.is_readonly() == True
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
try:
g.add_nodes(5)
fail = False
except DGLError:
fail = True
finally:
assert fail
g.readonly()
assert g._graph.is_readonly() == True
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
try:
g.add_nodes(5)
fail = False
except DGLError:
fail = True
finally:
assert fail
g.readonly(False)
assert g._graph.is_readonly() == False
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
try:
g.add_nodes(10)
g.add_edges([4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
[5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
fail = False
except DGLError:
fail = True
finally:
assert not fail
assert g.number_of_nodes() == 15
assert F.shape(g.ndata['x']) == (15, 3)
assert g.number_of_edges() == 14
assert F.shape(g.edata['x']) == (14, 4)
def test_find_edges():
g = dgl.DGLGraph()
g.add_nodes(10)
......@@ -388,19 +321,6 @@ def test_find_edges():
finally:
assert fail
g.readonly()
e = g.find_edges([1, 3, 2, 4])
assert F.asnumpy(e[0][0]) == 1 and F.asnumpy(e[0][1]) == 3 and F.asnumpy(e[0][2]) == 2 and F.asnumpy(e[0][3]) == 4
assert F.asnumpy(e[1][0]) == 2 and F.asnumpy(e[1][1]) == 4 and F.asnumpy(e[1][2]) == 3 and F.asnumpy(e[1][3]) == 5
try:
g.find_edges([10])
fail = False
except DGLError:
fail = True
finally:
assert fail
def test_ismultigraph():
g = dgl.DGLGraph()
g.add_nodes(10)
......@@ -417,5 +337,4 @@ if __name__ == '__main__':
test_mutation()
test_scipy_adjmat()
test_incmat()
test_readonly()
test_find_edges()
tests/compute/test_graph_index.py deleted 100644 → 0
View file @ 015acfd2
import os
import backend as F
import networkx as nx
import numpy as np
import scipy as sp
from scipy import sparse as spsp
import dgl
from dgl.graph_index import map_to_subgraph_nid, GraphIndex, create_graph_index
from dgl import utils
def generate_from_networkx():
edges = [[2, 3], [2, 5], [3, 0], [1, 0], [4, 3], [4, 5]]
nx_graph = nx.DiGraph()
nx_graph.add_edges_from(edges)
g = create_graph_index(nx_graph, readonly=False)
ig = create_graph_index(nx_graph, readonly=True)
return g, ig
def generate_from_edgelist():
edges = [[2, 3], [2, 5], [3, 0], [6, 10], [10, 3], [10, 15]]
g = create_graph_index(edges, readonly=False)
ig = create_graph_index(edges, readonly=True)
return g, ig
def generate_rand_graph(n):
arr = (sp.sparse.random(n, n, density=0.1, format='coo') != 0).astype(np.int64)
g = create_graph_index(arr, readonly=False)
ig = create_graph_index(arr, readonly=True)
return g, ig
def check_graph_equal(g1, g2):
adj1 = g1.adjacency_matrix(False, F.cpu())[0]
adj2 = g2.adjacency_matrix(False, F.cpu())[0]
assert np.all(F.asnumpy(adj1) == F.asnumpy(adj2))
def test_graph_gen():
g, ig = generate_from_edgelist()
check_graph_equal(g, ig)
g, ig = generate_rand_graph(10)
check_graph_equal(g, ig)
def sort_edges(edges):
edges = [e.tousertensor() for e in edges]
if np.prod(edges[2].shape) > 0:
val, idx = F.sort_1d(edges[2])
return (F.gather_row(edges[0], idx), F.gather_row(edges[1], idx), F.gather_row(edges[2], idx))
else:
return (edges[0], edges[1], edges[2])
def check_basics(g, ig):
assert g.number_of_nodes() == ig.number_of_nodes()
assert g.number_of_edges() == ig.number_of_edges()
edges = g.edges("srcdst")
iedges = ig.edges("srcdst")
assert F.array_equal(edges[0].tousertensor(), iedges[0].tousertensor())
assert F.array_equal(edges[1].tousertensor(), iedges[1].tousertensor())
assert F.array_equal(edges[2].tousertensor(), iedges[2].tousertensor())
edges = g.edges("eid")
iedges = ig.edges("eid")
assert F.array_equal(edges[0].tousertensor(), iedges[0].tousertensor())
assert F.array_equal(edges[1].tousertensor(), iedges[1].tousertensor())
assert F.array_equal(edges[2].tousertensor(), iedges[2].tousertensor())
for i in range(g.number_of_nodes()):
assert g.has_node(i) == ig.has_node(i)
for i in range(g.number_of_nodes()):
assert F.array_equal(g.predecessors(i).tousertensor(), ig.predecessors(i).tousertensor())
assert F.array_equal(g.successors(i).tousertensor(), ig.successors(i).tousertensor())
randv = np.random.randint(0, g.number_of_nodes(), 10)
randv = utils.toindex(randv)
in_src1, in_dst1, in_eids1 = sort_edges(g.in_edges(randv))
in_src2, in_dst2, in_eids2 = sort_edges(ig.in_edges(randv))
nnz = in_src2.shape[0]
assert F.array_equal(in_src1, in_src2)
assert F.array_equal(in_dst1, in_dst2)
assert F.array_equal(in_eids1, in_eids2)
out_src1, out_dst1, out_eids1 = sort_edges(g.out_edges(randv))
out_src2, out_dst2, out_eids2 = sort_edges(ig.out_edges(randv))
nnz = out_dst2.shape[0]
assert F.array_equal(out_dst1, out_dst2)
assert F.array_equal(out_src1, out_src2)
assert F.array_equal(out_eids1, out_eids2)
num_v = len(randv)
assert F.array_equal(g.in_degrees(randv).tousertensor(), ig.in_degrees(randv).tousertensor())
assert F.array_equal(g.out_degrees(randv).tousertensor(), ig.out_degrees(randv).tousertensor())
randv = randv.tousertensor()
for v in F.asnumpy(randv):
assert g.in_degree(v) == ig.in_degree(v)
assert g.out_degree(v) == ig.out_degree(v)
for u in F.asnumpy(randv):
for v in F.asnumpy(randv):
if len(g.edge_id(u, v)) == 1:
assert g.edge_id(u, v).tonumpy() == ig.edge_id(u, v).tonumpy()
assert g.has_edge_between(u, v) == ig.has_edge_between(u, v)
randv = utils.toindex(randv)
ids = g.edge_ids(randv, randv)[2].tonumpy()
assert sum(ig.edge_ids(randv, randv)[2].tonumpy() == ids, 0) == len(ids)
assert sum(g.has_edges_between(randv, randv).tonumpy() == ig.has_edges_between(randv, randv).tonumpy(), 0) == len(randv)
def test_basics():
g, ig = generate_from_edgelist()
check_basics(g, ig)
g, ig = generate_from_networkx()
check_basics(g, ig)
g, ig = generate_rand_graph(100)
check_basics(g, ig)
def test_node_subgraph():
num_vertices = 100
g, ig = generate_rand_graph(num_vertices)
# node_subgraph
randv1 = np.random.randint(0, num_vertices, 20)
randv = np.unique(randv1)
subg = g.node_subgraph(utils.toindex(randv))
subig = ig.node_subgraph(utils.toindex(randv))
check_basics(subg.graph, subig.graph)
check_graph_equal(subg.graph, subig.graph)
assert F.asnumpy(map_to_subgraph_nid(subg.induced_nodes, utils.toindex(randv1[0:10])).tousertensor()
== map_to_subgraph_nid(subig.induced_nodes, utils.toindex(randv1[0:10])).tousertensor()).sum(0).item() == 10
# node_subgraphs
randvs = []
subgs = []
for i in range(4):
randv = np.unique(np.random.randint(0, num_vertices, 20))
randvs.append(utils.toindex(randv))
subgs.append(g.node_subgraph(utils.toindex(randv)))
subigs= ig.node_subgraphs(randvs)
for i in range(4):
check_basics(subg.graph, subig.graph)
check_graph_equal(subgs[i].graph, subigs[i].graph)
def test_create_graph():
elist = [(1, 2), (0, 1), (0, 2)]
ig = dgl.DGLGraph(elist, readonly=True)
g = dgl.DGLGraph(elist, readonly=False)
for edge in elist:
assert g.edge_id(edge[0], edge[1]) == ig.edge_id(edge[0], edge[1])
data = [1, 2, 3]
rows = [1, 0, 0]
cols = [2, 1, 2]
mat = sp.sparse.coo_matrix((data, (rows, cols)))
g = dgl.DGLGraph(mat, readonly=False)
ig = dgl.DGLGraph(mat, readonly=True)
for edge in elist:
assert g.edge_id(edge[0], edge[1]) == ig.edge_id(edge[0], edge[1])
def test_load_csr():
n = 100
csr = (sp.sparse.random(n, n, density=0.1, format='csr') != 0).astype(np.int64)
# Load CSR normally.
idx = dgl.graph_index.from_csr(
utils.toindex(csr.indptr), utils.toindex(csr.indices), 'out')
assert idx.number_of_nodes() == n
assert idx.number_of_edges() == csr.nnz
src, dst, eid = idx.edges()
src, dst, eid = src.tousertensor(), dst.tousertensor(), eid.tousertensor()
coo = csr.tocoo()
assert np.all(F.asnumpy(src) == coo.row)
assert np.all(F.asnumpy(dst) == coo.col)
def test_edge_ids():
np.random.seed(0)
csr = (spsp.random(20, 20, density=0.1, format='csr') != 0).astype(np.int64)
#csr = csr.transpose()
g = dgl.DGLGraph(csr, readonly=True)
num_nodes = g.number_of_nodes()
in_edges = g._graph.in_edges(v=dgl.utils.toindex([2]))
src, dst, eids = g._graph.edge_ids(dgl.utils.toindex(in_edges[0]),
dgl.utils.toindex(in_edges[1]))
assert np.all(in_edges[0].tonumpy() == src.tonumpy())
assert np.all(in_edges[1].tonumpy() == dst.tonumpy())
if __name__ == '__main__':
test_basics()
test_edge_ids()
test_graph_gen()
test_node_subgraph()
test_create_graph()
test_load_csr()
tests/compute/test_hetero_basics.py
View file @ 44089c8b
......@@ -28,7 +28,7 @@ def reduce_func(nodes):
def apply_node_func(nodes):
return {'h' : nodes.data['h'] + nodes.data['accum']}
def generate_graph(index_dtype='int64', grad=False):
def generate_graph(idtype=F.int64, grad=False):
'''
s, d, eid
0, 1, 0
......@@ -49,8 +49,10 @@ def generate_graph(index_dtype='int64', grad=False):
8, 9, 15
9, 0, 16
'''
g = dgl.graph([(0,1), (1,9), (0,2), (2,9), (0,3), (3,9), (0,4), (4,9),
(0,5), (5,9), (0,6), (6,9), (0,7), (7,9), (0,8), (8,9), (9,0)], index_dtype=index_dtype)
u = F.tensor([0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8, 9])
v = F.tensor([1, 9, 2, 9, 3, 9, 4, 9, 5, 9, 6, 9, 7, 9, 8, 9, 0])
g = dgl.graph((u, v), idtype=idtype)
assert g.device == F.ctx()
ncol = F.randn((10, D))
ecol = F.randn((17, D))
if grad:
......@@ -65,33 +67,32 @@ def generate_graph(index_dtype='int64', grad=False):
@parametrize_dtype
def test_isolated_nodes(index_dtype):
g = dgl.graph([(0, 1), (1, 2)], num_nodes=5, index_dtype=index_dtype)
assert g._idtype_str == index_dtype
def test_isolated_nodes(idtype):
g = dgl.graph([(0, 1), (1, 2)], num_nodes=5, idtype=idtype, device=F.ctx())
assert g.number_of_nodes() == 5
# Test backward compatibility
g = dgl.graph([(0, 1), (1, 2)], card=5, index_dtype=index_dtype)
g = dgl.graph([(0, 1), (1, 2)], card=5, idtype=idtype, device=F.ctx())
assert g.number_of_nodes() == 5
g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays',
'game', num_nodes=(5, 7), index_dtype=index_dtype)
assert g._idtype_str == index_dtype
'game', num_nodes=(5, 7), idtype=idtype, device=F.ctx())
assert g.idtype == idtype
assert g.number_of_nodes('user') == 5
assert g.number_of_nodes('game') == 7
# Test backward compatibility
g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays',
'game', card=(5, 7), index_dtype=index_dtype)
assert g._idtype_str == index_dtype
'game', card=(5, 7), idtype=idtype, device=F.ctx())
assert g.idtype == idtype
assert g.number_of_nodes('user') == 5
assert g.number_of_nodes('game') == 7
@parametrize_dtype
def test_batch_setter_getter(index_dtype):
def test_batch_setter_getter(idtype):
def _pfc(x):
return list(F.zerocopy_to_numpy(x)[:,0])
g = generate_graph(index_dtype)
g = generate_graph(idtype)
# set all nodes
g.ndata['h'] = F.zeros((10, D))
assert F.allclose(g.ndata['h'], F.zeros((10, D)))
......@@ -101,11 +102,11 @@ def test_batch_setter_getter(index_dtype):
assert len(g.ndata) == old_len - 1
g.ndata['h'] = F.zeros((10, D))
# set partial nodes
u = F.tensor([1, 3, 5], F.data_type_dict[index_dtype])
u = F.tensor([1, 3, 5], idtype)
g.nodes[u].data['h'] = F.ones((3, D))
assert _pfc(g.ndata['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
# get partial nodes
u = F.tensor([1, 2, 3], F.data_type_dict[index_dtype])
u = F.tensor([1, 2, 3], idtype)
assert _pfc(g.nodes[u].data['h']) == [1., 0., 1.]
'''
......@@ -137,44 +138,44 @@ def test_batch_setter_getter(index_dtype):
assert len(g.edata) == old_len - 1
g.edata['l'] = F.zeros((17, D))
# set partial edges (many-many)
u = F.tensor([0, 0, 2, 5, 9], dtype=F.data_type_dict[index_dtype])
v = F.tensor([1, 3, 9, 9, 0], dtype=F.data_type_dict[index_dtype])
u = F.tensor([0, 0, 2, 5, 9], dtype=idtype)
v = F.tensor([1, 3, 9, 9, 0], dtype=idtype)
g.edges[u, v].data['l'] = F.ones((5, D))
truth = [0.] * 17
truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
assert _pfc(g.edata['l']) == truth
# set partial edges (many-one)
u = F.tensor([3, 4, 6], dtype=F.data_type_dict[index_dtype])
v = F.tensor([9], dtype=F.data_type_dict[index_dtype])
u = F.tensor([3, 4, 6], dtype=idtype)
v = F.tensor([9], dtype=idtype)
g.edges[u, v].data['l'] = F.ones((3, D))
truth[5] = truth[7] = truth[11] = 1.
assert _pfc(g.edata['l']) == truth
# set partial edges (one-many)
u = F.tensor([0], dtype=F.data_type_dict[index_dtype])
v = F.tensor([4, 5, 6], dtype=F.data_type_dict[index_dtype])
u = F.tensor([0], dtype=idtype)
v = F.tensor([4, 5, 6], dtype=idtype)
g.edges[u, v].data['l'] = F.ones((3, D))
truth[6] = truth[8] = truth[10] = 1.
assert _pfc(g.edata['l']) == truth
# get partial edges (many-many)
u = F.tensor([0, 6, 0], dtype=F.data_type_dict[index_dtype])
v = F.tensor([6, 9, 7], dtype=F.data_type_dict[index_dtype])
u = F.tensor([0, 6, 0], dtype=idtype)
v = F.tensor([6, 9, 7], dtype=idtype)
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
# get partial edges (many-one)
u = F.tensor([5, 6, 7], dtype=F.data_type_dict[index_dtype])
v = F.tensor([9], dtype=F.data_type_dict[index_dtype])
u = F.tensor([5, 6, 7], dtype=idtype)
v = F.tensor([9], dtype=idtype)
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
# get partial edges (one-many)
u = F.tensor([0], dtype=F.data_type_dict[index_dtype])
v = F.tensor([3, 4, 5], dtype=F.data_type_dict[index_dtype])
u = F.tensor([0], dtype=idtype)
v = F.tensor([3, 4, 5], dtype=idtype)
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 1.]
@parametrize_dtype
def test_batch_setter_autograd(index_dtype):
g = generate_graph(index_dtype=index_dtype, grad=True)
def test_batch_setter_autograd(idtype):
g = generate_graph(idtype=idtype, grad=True)
h1 = g.ndata['h']
# partial set
v = F.tensor([1, 2, 8], F.data_type_dict[index_dtype])
v = F.tensor([1, 2, 8], idtype)
hh = F.attach_grad(F.zeros((len(v), D)))
with F.record_grad():
g.nodes[v].data['h'] = hh
......@@ -185,7 +186,7 @@ def test_batch_setter_autograd(index_dtype):
@parametrize_dtype
def atest_nx_conversion(index_dtype):
def test_nx_conversion(idtype):
# check conversion between networkx and DGLGraph
def _check_nx_feature(nxg, nf, ef):
......@@ -222,20 +223,22 @@ def atest_nx_conversion(index_dtype):
n3 = F.randn((5, 4))
e1 = F.randn((4, 5))
e2 = F.randn((4, 7))
g = dgl.graph([(0,2),(1,4),(3,0),(4,3)], index_dtype=index_dtype)
g = dgl.graph([(0,2),(1,4),(3,0),(4,3)], idtype=idtype, device=F.ctx())
g.ndata.update({'n1': n1, 'n2': n2, 'n3': n3})
g.edata.update({'e1': e1, 'e2': e2})
# convert to networkx
nxg = dgl.to_networkx(g, node_attrs=['n1', 'n3'], edge_attrs=['e1', 'e2'])
nxg = dgl.to_networkx(g.cpu(), node_attrs=['n1', 'n3'], edge_attrs=['e1', 'e2'])
assert len(nxg) == 5
assert nxg.size() == 4
_check_nx_feature(nxg, {'n1': n1, 'n3': n3}, {'e1': e1, 'e2': e2})
# convert to DGLGraph, nx graph has id in edge feature
# use id feature to test non-tensor copy
g = dgl.graph(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'], index_dtype=index_dtype)
assert g._idtype_str == index_dtype
g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'], idtype=idtype)
assert g.idtype == idtype
assert g.device == F.cpu()
g = g.to(F.ctx())
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
......@@ -247,7 +250,7 @@ def atest_nx_conversion(index_dtype):
assert F.allclose(g.ndata['n1'], n1)
# with id in nx edge feature, e1 should follow original order
assert F.allclose(g.edata['e1'], e1)
assert F.array_equal(g.edata['id'], F.copy_to(F.arange(0, 4), F.cpu()))
assert F.array_equal(g.edata['id'], F.arange(0, 4, F.dtype(g.edata['id'])))
# test conversion after modifying DGLGraph
# TODO(minjie): enable after mutation is supported
......@@ -270,7 +273,7 @@ def atest_nx_conversion(index_dtype):
for _, _, attr in nxg.edges(data=True):
attr.pop('id')
# test with a new graph
g = dgl.graph(nxg , node_attrs=['n1'], edge_attrs=['e1'])
g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1'], idtype=idtype)
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
......@@ -286,86 +289,35 @@ def atest_nx_conversion(index_dtype):
edge_feat = F.cat(edge_feat, 0)
assert F.allclose(g.edata['e1'], edge_feat)
# Test converting from a networkx graph whose nodes are
# not labeled with consecutive-integers.
nxg = nx.cycle_graph(5)
nxg.remove_nodes_from([0, 4])
for u in nxg.nodes():
nxg.nodes[u]['h'] = F.tensor([u])
for u, v, d in nxg.edges(data=True):
d['h'] = F.tensor([u, v])
g = dgl.DGLGraph()
g.from_networkx(nxg, node_attrs=['h'], edge_attrs=['h'])
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 4
assert g.has_edge_between(0, 1)
assert g.has_edge_between(1, 2)
assert F.allclose(g.ndata['h'], F.tensor([[1.], [2.], [3.]]))
assert F.allclose(g.edata['h'], F.tensor([[1., 2.], [1., 2.],
[2., 3.], [2., 3.]]))
@parametrize_dtype
def test_batch_send(index_dtype):
g = generate_graph(index_dtype=index_dtype)
def _fmsg(edges):
assert tuple(F.shape(edges.src['h'])) == (5, D)
return {'m' : edges.src['h']}
# many-many send
u = F.tensor([0, 0, 0, 0, 0], dtype=F.data_type_dict[index_dtype])
v = F.tensor([1, 2, 3, 4, 5], dtype=F.data_type_dict[index_dtype])
g.send((u, v), _fmsg)
# one-many send
u = F.tensor([0], dtype=F.data_type_dict[index_dtype])
v = F.tensor([1, 2, 3, 4, 5], dtype=F.data_type_dict[index_dtype])
g.send((u, v), _fmsg)
# many-one send
u = F.tensor([1, 2, 3, 4, 5], dtype=F.data_type_dict[index_dtype])
v = F.tensor([9], dtype=F.data_type_dict[index_dtype])
g.send((u, v), _fmsg)
@parametrize_dtype
def test_batch_recv(index_dtype):
# basic recv test
g = generate_graph(index_dtype=index_dtype)
u = F.tensor([0, 0, 0, 4, 5, 6], dtype=F.data_type_dict[index_dtype])
v = F.tensor([1, 2, 3, 9, 9, 9], dtype=F.data_type_dict[index_dtype])
reduce_msg_shapes.clear()
g.send((u, v), message_func)
g.recv(F.astype(F.unique(v), F.data_type_dict[index_dtype]), reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
@parametrize_dtype
def test_apply_nodes(index_dtype):
def test_apply_nodes(idtype):
def _upd(nodes):
return {'h' : nodes.data['h'] * 2}
g = generate_graph(index_dtype=index_dtype)
g = generate_graph(idtype=idtype)
old = g.ndata['h']
g.apply_nodes(_upd)
assert F.allclose(old * 2, g.ndata['h'])
u = F.tensor([0, 3, 4, 6], F.data_type_dict[index_dtype])
u = F.tensor([0, 3, 4, 6], idtype)
g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
assert F.allclose(F.gather_row(g.ndata['h'], u), F.zeros((4, D)))
@parametrize_dtype
def test_apply_edges(index_dtype):
def test_apply_edges(idtype):
def _upd(edges):
return {'w' : edges.data['w'] * 2}
g = generate_graph(index_dtype=index_dtype)
g = generate_graph(idtype=idtype)
old = g.edata['w']
g.apply_edges(_upd)
assert F.allclose(old * 2, g.edata['w'])
u = F.tensor([0, 0, 0, 4, 5, 6], F.data_type_dict[index_dtype])
v = F.tensor([1, 2, 3, 9, 9, 9], F.data_type_dict[index_dtype])
u = F.tensor([0, 0, 0, 4, 5, 6], idtype)
v = F.tensor([1, 2, 3, 9, 9, 9], idtype)
g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
eid = F.tensor(g.edge_ids(u, v), F.data_type_dict[index_dtype])
eid = F.tensor(g.edge_ids(u, v), idtype)
assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
@parametrize_dtype
def test_update_routines(index_dtype):
g = generate_graph(index_dtype=index_dtype)
def test_update_routines(idtype):
g = generate_graph(idtype=idtype)
# send_and_recv
reduce_msg_shapes.clear()
......@@ -381,14 +333,14 @@ def test_update_routines(index_dtype):
pass
# pull
v = F.tensor([1, 2, 3, 9], F.data_type_dict[index_dtype])
v = F.tensor([1, 2, 3, 9], idtype)
reduce_msg_shapes.clear()
g.pull(v, message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
reduce_msg_shapes.clear()
# push
v = F.tensor([0, 1, 2, 3], F.data_type_dict[index_dtype])
v = F.tensor([0, 1, 2, 3], idtype)
reduce_msg_shapes.clear()
g.push(v, message_func, reduce_func, apply_node_func)
assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
......@@ -401,81 +353,9 @@ def test_update_routines(index_dtype):
reduce_msg_shapes.clear()
@parametrize_dtype
def test_recv_0deg(index_dtype):
# test recv with 0deg nodes;
g = dgl.graph([(0,1)], index_dtype=index_dtype)
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h' : nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype, ctx)
g.set_n_initializer(_init2, 'h')
# test#1: recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1), _message)
g.recv([0, 1], _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# test#2: recv only 0deg node is equal to apply
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1), _message)
g.recv(0, _reduce, _apply)
new = g.ndata.pop('h')
# 0deg check: equal to apply_nodes
assert F.allclose(new[0], 2 * old[0])
# non-0deg check: untouched
assert F.allclose(new[1], old[1])
@parametrize_dtype
def test_recv_0deg_newfld(index_dtype):
# test recv with 0deg nodes; the reducer also creates a new field
g = dgl.graph([(0,1)], index_dtype=index_dtype)
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
return {'h1' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h1' : nodes.data['h1'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype=dtype, ctx=ctx)
# test#1: recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.set_n_initializer(_init2, 'h1')
g.ndata['h'] = old
g.send((0, 1), _message)
g.recv([0, 1], _reduce, _apply)
new = g.ndata.pop('h1')
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# test#2: recv only 0deg node
old = F.randn((2, 5))
g.ndata['h'] = old
g.ndata['h1'] = F.full((2, 5), -1, F.int64) # this is necessary
g.send((0, 1), _message)
g.recv(0, _reduce, _apply)
new = g.ndata.pop('h1')
# 0deg check: fallback to apply
assert F.allclose(new[0], F.full_1d(5, -2, F.int64))
# non-0deg check: not changed
assert F.allclose(new[1], F.full_1d(5, -1, F.int64))
@parametrize_dtype
def test_update_all_0deg(index_dtype):
def test_update_all_0deg(idtype):
# test#1
g = dgl.graph([(1,0), (2,0), (3,0), (4,0)], index_dtype=index_dtype)
g = dgl.graph([(1,0), (2,0), (3,0), (4,0)], idtype=idtype, device=F.ctx())
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
......@@ -496,7 +376,7 @@ def test_update_all_0deg(index_dtype):
assert F.allclose(new_repr[0], 2 * F.sum(old_repr, 0))
# test#2:
g = dgl.graph([], num_nodes=5, index_dtype=index_dtype)
g = dgl.graph([], num_nodes=5, idtype=idtype, device=F.ctx())
g.set_n_initializer(_init2, 'h')
g.ndata['h'] = old_repr
g.update_all(_message, _reduce, _apply)
......@@ -505,8 +385,8 @@ def test_update_all_0deg(index_dtype):
assert F.allclose(new_repr, 2*old_repr)
@parametrize_dtype
def test_pull_0deg(index_dtype):
g = dgl.graph([(0,1)], index_dtype=index_dtype)
def test_pull_0deg(idtype):
g = dgl.graph([(0,1)], idtype=idtype, device=F.ctx())
def _message(edges):
return {'m' : edges.src['h']}
def _reduce(nodes):
......@@ -537,8 +417,8 @@ def test_pull_0deg(index_dtype):
assert F.allclose(new[1], old[1])
@parametrize_dtype
def test_send_multigraph(index_dtype):
g = dgl.graph([(0,1), (0,1), (0,1), (2,1)], index_dtype=index_dtype)
def test_send_multigraph(idtype):
g = dgl.graph([(0,1), (0,1), (0,1), (2,1)], idtype=idtype, device=F.ctx())
def _message_a(edges):
return {'a': edges.data['a']}
......@@ -554,46 +434,6 @@ def test_send_multigraph(index_dtype):
# send by eid
old_repr = F.randn((4, 5))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2]))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2, 3], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
# send on multigraph
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(([0, 2], [1, 1]), _message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], F.max(old_repr, 0))
# consecutive send and send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send((2, 1), _message_a)
g.send([0, 1], message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0] * 3, old_repr[1] * 3, old_repr[3]))
# consecutive send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(0, message_func=_message_a)
g.send(1, message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[1] * 3))
# send_and_recv_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
......@@ -636,8 +476,8 @@ def _test_dynamic_addition():
assert len(g.edata['h1']) == len(g.edata['h2'])
@parametrize_dtype
def test_repr(index_dtype):
G = dgl.graph([(0,1), (0,2), (1,2)], num_nodes=10, index_dtype=index_dtype)
def test_repr(idtype):
G = dgl.graph([(0,1), (0,2), (1,2)], num_nodes=10, idtype=idtype, device=F.ctx())
repr_string = G.__repr__()
print(repr_string)
G.ndata['x'] = F.zeros((10, 5))
......@@ -647,7 +487,7 @@ def test_repr(index_dtype):
@parametrize_dtype
def test_group_apply_edges(index_dtype):
def test_group_apply_edges(idtype):
def edge_udf(edges):
h = F.sum(edges.data['feat'] * (edges.src['h'] + edges.dst['h']), dim=2)
normalized_feat = F.softmax(h, dim=1)
......@@ -660,7 +500,7 @@ def test_group_apply_edges(index_dtype):
elist.append((1, v))
for v in [2, 3, 4, 5, 6, 7, 8]:
elist.append((2, v))
g = dgl.graph(elist, index_dtype=index_dtype)
g = dgl.graph(elist, idtype=idtype, device=F.ctx())
g.ndata['h'] = F.randn((g.number_of_nodes(), D))
g.edata['feat'] = F.randn((g.number_of_edges(), D))
......@@ -683,8 +523,9 @@ def test_group_apply_edges(index_dtype):
_test('dst')
@parametrize_dtype
def test_local_var(index_dtype):
g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], index_dtype=index_dtype)
def test_local_var(idtype):
g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], idtype=idtype, device=F.ctx())
g = g.to(F.ctx())
g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
g.edata['w'] = F.zeros((g.number_of_edges(), 4))
# test override
......@@ -721,7 +562,8 @@ def test_local_var(index_dtype):
assert 'ww' not in g.edata
# test initializer1
g = dgl.graph([(0,1), (1,1)])
g = dgl.graph([(0,1), (1,1)], idtype=idtype, device=F.ctx())
g = g.to(F.ctx())
g.set_n_initializer(dgl.init.zero_initializer)
def foo(g):
g = g.local_var()
......@@ -741,8 +583,8 @@ def test_local_var(index_dtype):
foo(g)
@parametrize_dtype
def test_local_scope(index_dtype):
g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], index_dtype=index_dtype)
def test_local_scope(idtype):
g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], idtype=idtype, device=F.ctx())
g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
g.edata['w'] = F.zeros((g.number_of_edges(), 4))
# test override
......@@ -793,7 +635,7 @@ def test_local_scope(index_dtype):
assert 'ww' not in g.edata
# test initializer1
g = dgl.graph([(0,1), (1,1)], index_dtype=index_dtype)
g = dgl.graph([(0,1), (1,1)], idtype=idtype, device=F.ctx())
g.set_n_initializer(dgl.init.zero_initializer)
def foo(g):
with g.local_scope():
......@@ -813,26 +655,25 @@ def test_local_scope(index_dtype):
foo(g)
@parametrize_dtype
def test_issue_1088(index_dtype):
def test_issue_1088(idtype):
# This test ensures that message passing on a heterograph with one edge type
# would not crash (GitHub issue #1088).
import dgl.function as fn
g = dgl.heterograph({('U', 'E', 'V'): ([0, 1, 2], [1, 2, 3])}, index_dtype=index_dtype)
g = dgl.heterograph({('U', 'E', 'V'): ([0, 1, 2], [1, 2, 3])}, idtype=idtype, device=F.ctx())
g.nodes['U'].data['x'] = F.randn((3, 3))
g.update_all(fn.copy_u('x', 'm'), fn.sum('m', 'y'))
if __name__ == '__main__':
# test_isolated_nodes("int32")
# test_nx_conversion()
# test_batch_setter_getter("int32")
#test_isolated_nodes("int32")
test_batch_setter_getter(F.int32)
# test_batch_recv("int64")
test_apply_edges("int32")
# test_apply_edges("int32")
# test_batch_setter_autograd()
# test_batch_send()
# test_batch_recv()
# test_apply_nodes()
# test_apply_edges()
# test_update_routines()
#test_update_routines(F.int32)
# test_recv_0deg()
# test_recv_0deg_newfld()
# test_update_all_0deg()
......@@ -843,4 +684,5 @@ if __name__ == '__main__':
# test_group_apply_edges()
# test_local_var()
# test_local_scope()
test_issue_1088('int64')
#test_issue_1088('int64')
pass
tests/compute/test_heterograph.py
View file @ 44089c8b
......@@ -8,10 +8,11 @@ import backend as F
import networkx as nx
import unittest, pytest
from dgl import DGLError
from dgl.heterograph_index import joint_union
from utils import parametrize_dtype
import test_utils
from test_utils import parametrize_dtype, get_cases
from scipy.sparse import rand
def create_test_heterograph(index_dtype):
def create_test_heterograph(idtype):
# test heterograph from the docstring, plus a user -- wishes -- game relation
# 3 users, 2 games, 2 developers
# metagraph:
......@@ -27,19 +28,20 @@ def create_test_heterograph(index_dtype):
wishes_nx.add_edge('u0', 'g1', id=0)
wishes_nx.add_edge('u2', 'g0', id=1)
follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows', index_dtype=index_dtype)
plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game', index_dtype=index_dtype)
wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game', index_dtype=index_dtype)
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game', index_dtype=index_dtype)
assert follows_g._idtype_str == index_dtype
assert plays_g._idtype_str == index_dtype
assert wishes_g._idtype_str == index_dtype
assert develops_g._idtype_str == index_dtype
follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows', idtype=idtype, device=F.ctx())
plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game', idtype=idtype, device=F.ctx())
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game', idtype=idtype, device=F.ctx())
assert follows_g.idtype == idtype
assert plays_g.idtype == idtype
assert wishes_g.idtype == idtype
assert develops_g.idtype == idtype
g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
assert g._idtype_str == index_dtype
assert g.idtype == idtype
assert g.device == F.ctx()
return g
def create_test_heterograph1(index_dtype):
def create_test_heterograph1(idtype):
edges = []
edges.extend([(0,1), (1,2)]) # follows
edges.extend([(0,3), (1,3), (2,4), (1,4)]) # plays
......@@ -47,29 +49,31 @@ def create_test_heterograph1(index_dtype):
edges.extend([(5,3), (6,4)]) # develops
ntypes = F.tensor([0, 0, 0, 1, 1, 2, 2])
etypes = F.tensor([0, 0, 1, 1, 1, 1, 2, 2, 3, 3])
g0 = dgl.graph(edges)
g0 = dgl.graph(edges, idtype=idtype, device=F.ctx())
g0.ndata[dgl.NTYPE] = ntypes
g0.edata[dgl.ETYPE] = etypes
return dgl.to_hetero(g0, ['user', 'game', 'developer'], ['follows', 'plays', 'wishes', 'develops'])
def create_test_heterograph2(index_dtype):
def create_test_heterograph2(idtype):
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
wishes_nx = nx.DiGraph()
wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
wishes_nx.add_edge('u0', 'g1', id=0)
wishes_nx.add_edge('u2', 'g0', id=1)
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')
g = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): plays_spmat,
('user', 'wishes', 'game'): wishes_nx,
('developer', 'develops', 'game'): develops_g,
})
('developer', 'develops', 'game'): (F.tensor([0, 1]), F.tensor([0, 1])),
}, idtype=idtype, device=F.ctx())
assert g.idtype == idtype
assert g.device == F.ctx()
return g
def create_test_heterograph3(index_dtype):
def create_test_heterograph3(idtype):
device = F.ctx()
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
wishes_nx = nx.DiGraph()
wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
......@@ -77,82 +81,134 @@ def create_test_heterograph3(index_dtype):
wishes_nx.add_edge('u0', 'g1', id=0)
wishes_nx.add_edge('u2', 'g0', id=1)
follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows', _restrict_format='coo')
plays_g = dgl.bipartite(
[(0, 0), (1, 0), (2, 1), (1, 1)], 'user', 'plays', 'game', _restrict_format='coo')
wishes_g = dgl.bipartite([(0, 1), (2, 0)], 'user', 'wishes', 'game', _restrict_format='coo')
develops_g = dgl.bipartite(
[(0, 0), (1, 1)], 'developer', 'develops', 'game', _restrict_format='coo')
follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows',
idtype=idtype, device=device).formats('coo')
plays_g = dgl.bipartite([(0, 0), (1, 0), (2, 1), (1, 1)], 'user', 'plays', 'game',
idtype=idtype, device=device).formats('coo')
wishes_g = dgl.bipartite([(0, 1), (2, 0)], 'user', 'wishes', 'game',
idtype=idtype, device=device).formats('coo')
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game',
idtype=idtype, device=device).formats('coo')
g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
assert g.idtype == idtype
assert g.device == device
return g
def create_test_heterograph4(idtype):
g = dgl.heterograph({
('user', 'plays', 'game'): (F.tensor([0, 1, 1, 2], dtype=idtype),
F.tensor([0, 0, 1, 1], dtype=idtype)),
('developer', 'develops', 'game'): (F.tensor([0, 1], dtype=idtype),
F.tensor([0, 1], dtype=idtype))},
idtype=idtype, device=F.ctx())
g.ndata['h'] = {'user' : F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()),
'game' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
'developer' : F.copy_to(F.tensor([3, 3], dtype=idtype), ctx=F.ctx())}
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 1, 1, 1], dtype=idtype), ctx=F.ctx())
return g
def create_test_heterograph5(idtype):
g = dgl.heterograph({
('user', 'follows', 'user'): (F.tensor([0, 1, 1, 2, 2, 2], dtype=idtype),
F.tensor([0, 0, 1, 1, 2, 2], dtype=idtype)),
('user', 'plays', 'game'): (F.tensor([0, 1], dtype=idtype),
F.tensor([0, 1], dtype=idtype))},
idtype=idtype, device=F.ctx())
g.ndata['h'] = {'user' : F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()),
'game' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g.edges['follows'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4, 5, 6], dtype=idtype), ctx=F.ctx())
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
return g
def create_test_heterograph6(idtype):
g = dgl.heterograph({
('user', 'follows', 'user'): (F.tensor([1, 2], dtype=idtype),
F.tensor([0, 1], dtype=idtype)),
('user', 'plays', 'game'): (F.tensor([0, 1], dtype=idtype),
F.tensor([0, 1], dtype=idtype))},
idtype=idtype, device=F.ctx())
g.ndata['h'] = {'user' : F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()),
'game' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g.edges['follows'].data['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
return g
def get_redfn(name):
return getattr(F, name)
@parametrize_dtype
def test_create(index_dtype):
g0 = create_test_heterograph(index_dtype)
g1 = create_test_heterograph1(index_dtype)
g2 = create_test_heterograph2(index_dtype)
def test_create(idtype):
device = F.ctx()
g0 = create_test_heterograph(idtype)
g1 = create_test_heterograph1(idtype)
g2 = create_test_heterograph2(idtype)
assert set(g0.ntypes) == set(g1.ntypes) == set(g2.ntypes)
assert set(g0.canonical_etypes) == set(g1.canonical_etypes) == set(g2.canonical_etypes)
# create from nx complete bipartite graph
nxg = nx.complete_bipartite_graph(3, 4)
g = dgl.bipartite(nxg, 'user', 'plays', 'game')
g = dgl.bipartite(nxg, 'user', 'plays', 'game', idtype=idtype, device=device)
assert g.ntypes == ['user', 'game']
assert g.etypes == ['plays']
assert g.number_of_edges() == 12
assert g.idtype == idtype
assert g.device == device
# create from scipy
spmat = ssp.coo_matrix(([1,1,1], ([0, 0, 1], [2, 3, 2])), shape=(4, 4))
g = dgl.graph(spmat)
g = dgl.graph(spmat, idtype=idtype, device=device)
assert g.number_of_nodes() == 4
assert g.number_of_edges() == 3
assert g.idtype == idtype
assert g.device == device
# test inferring number of nodes for heterograph
g = dgl.heterograph({
('l0', 'e0', 'l1'): [(0, 1), (0, 2)],
('l0', 'e1', 'l2'): [(2, 2)],
('l2', 'e2', 'l2'): [(1, 1), (3, 3)],
})
}, idtype=idtype, device=device)
assert g.number_of_nodes('l0') == 3
assert g.number_of_nodes('l1') == 3
assert g.number_of_nodes('l2') == 4
assert g.idtype == idtype
assert g.device == device
# test if validate flag works
# homo graph
fail = False
try:
with pytest.raises(DGLError):
g = dgl.graph(
([0, 0, 0, 1, 1, 2], [0, 1, 2, 0, 1, 2]),
num_nodes=2,
validate=True
validate=True,
idtype=idtype, device=device
)
except DGLError:
fail = True
finally:
assert fail, "should catch a DGLError because node ID is out of bound."
# bipartite graph
def _test_validate_bipartite(card):
fail = False
try:
with pytest.raises(DGLError):
g = dgl.bipartite(
([0, 0, 1, 1, 2], [1, 1, 2, 2, 3]),
num_nodes=card,
validate=True
validate=True,
idtype=idtype, device=device
)
except DGLError:
fail = True
finally:
assert fail, "should catch a DGLError because node ID is out of bound."
_test_validate_bipartite((3, 3))
_test_validate_bipartite((2, 4))
# test from_scipy
num_nodes = 10
density = 0.25
for fmt in ['csr', 'coo', 'csc']:
adj = rand(num_nodes, num_nodes, density=density, format=fmt)
g = dgl.from_scipy(adj, eweight_name='w', idtype=idtype)
assert g.idtype == idtype
assert g.device == F.cpu()
assert F.array_equal(g.edata['w'], F.copy_to(F.tensor(adj.data), F.cpu()))
@parametrize_dtype
def test_query(index_dtype):
g = create_test_heterograph(index_dtype)
def test_query(idtype):
g = create_test_heterograph(idtype)
ntypes = ['user', 'game', 'developer']
canonical_etypes = [
......@@ -196,17 +252,16 @@ def test_query(index_dtype):
F.asnumpy(g.has_nodes([0, n], ntype)).astype('int32'), [1, 0])
assert not g.is_multigraph
assert g.is_readonly
for etype in etypes:
srcs, dsts = edges[etype]
for src, dst in zip(srcs, dsts):
assert g.has_edge_between(src, dst, etype)
assert g.has_edges_between(src, dst, etype)
assert F.asnumpy(g.has_edges_between(srcs, dsts, etype)).all()
srcs, dsts = negative_edges[etype]
for src, dst in zip(srcs, dsts):
assert not g.has_edge_between(src, dst, etype)
assert not g.has_edges_between(src, dst, etype)
assert not F.asnumpy(g.has_edges_between(srcs, dsts, etype)).any()
srcs, dsts = edges[etype]
......@@ -218,7 +273,7 @@ def test_query(index_dtype):
u, v = g.in_edges([0], etype=etype)
assert F.asnumpy(v).tolist() == [0] * len(pred)
assert set(F.asnumpy(u).tolist()) == set(pred)
assert g.in_degree(0, etype) == len(pred)
assert g.in_degrees(0, etype) == len(pred)
# successors & out_edges & out_degree
succ = [d for s, d in zip(srcs, dsts) if s == 0]
......@@ -226,27 +281,28 @@ def test_query(index_dtype):
u, v = g.out_edges([0], etype=etype)
assert F.asnumpy(u).tolist() == [0] * len(succ)
assert set(F.asnumpy(v).tolist()) == set(succ)
assert g.out_degree(0, etype) == len(succ)
assert g.out_degrees(0, etype) == len(succ)
# edge_id & edge_ids
for i, (src, dst) in enumerate(zip(srcs, dsts)):
assert g.edge_id(src, dst, etype=etype) == i
assert F.asnumpy(g.edge_id(src, dst, etype=etype, return_array=True)).tolist() == [i]
assert g.edge_ids(src, dst, etype=etype) == i
_, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
assert eid == i
assert F.asnumpy(g.edge_ids(srcs, dsts, etype=etype)).tolist() == list(range(n_edges))
u, v, e = g.edge_ids(srcs, dsts, etype=etype, return_uv=True)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
assert u[e].tolist() == srcs
assert v[e].tolist() == dsts
# find_edges
for edge_ids in [list(range(n_edges)), np.arange(n_edges), F.astype(F.arange(0, n_edges), g.idtype)]:
u, v = g.find_edges(edge_ids, etype)
for eid in [list(range(n_edges)), np.arange(n_edges), F.astype(F.arange(0, n_edges), g.idtype)]:
u, v = g.find_edges(eid, etype)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
# all_edges.
for order in ['eid']:
u, v, e = g.all_edges('all', order, etype)
u, v, e = g.edges('all', order, etype)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
......@@ -275,12 +331,14 @@ def test_query(index_dtype):
'wishes': ([0, 1], [0, 1]),
'develops': ([0, 1], [1, 0]),
}
g = create_test_heterograph(index_dtype)
_test(g)
g = create_test_heterograph1(index_dtype)
g = create_test_heterograph(idtype)
_test(g)
g = create_test_heterograph3(index_dtype)
g = create_test_heterograph1(idtype)
_test(g)
if F._default_context_str != 'gpu':
# XXX: CUDA COO operators have not been live yet.
g = create_test_heterograph3(idtype)
_test(g)
etypes = canonical_etypes
edges = {
......@@ -296,16 +354,19 @@ def test_query(index_dtype):
('user', 'wishes', 'game'): ([0, 1], [0, 1]),
('developer', 'develops', 'game'): ([0, 1], [1, 0]),
}
g = create_test_heterograph(index_dtype)
g = create_test_heterograph(idtype)
_test(g)
g = create_test_heterograph1(index_dtype)
_test(g)
g = create_test_heterograph3(index_dtype)
g = create_test_heterograph1(idtype)
_test(g)
if F._default_context_str != 'gpu':
# XXX: CUDA COO operators have not been live yet.
g = create_test_heterograph3(idtype)
_test(g)
# test repr
print(g)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU does not have COO impl.")
def test_hypersparse():
N1 = 1 << 50 # should crash if allocated a CSR
N2 = 1 << 48
......@@ -313,19 +374,20 @@ def test_hypersparse():
g = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1)],
('user', 'plays', 'game'): [(0, N2)]},
{'user': N1, 'game': N1})
{'user': N1, 'game': N1},
idtype=F.int64, device=F.ctx())
assert g.number_of_nodes('user') == N1
assert g.number_of_nodes('game') == N1
assert g.number_of_edges('follows') == 1
assert g.number_of_edges('plays') == 1
assert g.has_edge_between(0, 1, 'follows')
assert not g.has_edge_between(0, 0, 'follows')
assert g.has_edges_between(0, 1, 'follows')
assert not g.has_edges_between(0, 0, 'follows')
mask = F.asnumpy(g.has_edges_between([0, 0], [0, 1], 'follows')).tolist()
assert mask == [0, 1]
assert g.has_edge_between(0, N2, 'plays')
assert not g.has_edge_between(0, 0, 'plays')
assert g.has_edges_between(0, N2, 'plays')
assert not g.has_edges_between(0, 0, 'plays')
mask = F.asnumpy(g.has_edges_between([0, 0], [0, N2], 'plays')).tolist()
assert mask == [0, 1]
......@@ -339,10 +401,8 @@ def test_hypersparse():
assert F.asnumpy(g.predecessors(N2, 'plays')).tolist() == [0]
assert F.asnumpy(g.successors(N2, 'plays')).tolist() == []
assert g.edge_id(0, 1, etype='follows') == 0
assert g.edge_id(0, N2, etype='plays') == 0
assert F.asnumpy(g.edge_ids(0, 1, etype='follows')).tolist() == [0]
assert F.asnumpy(g.edge_ids(0, N2, etype='plays')).tolist() == [0]
assert g.edge_ids(0, 1, etype='follows') == 0
assert g.edge_ids(0, N2, etype='plays') == 0
u, v = g.find_edges([0], 'follows')
assert F.asnumpy(u).tolist() == [0]
......@@ -359,17 +419,17 @@ def test_hypersparse():
assert F.asnumpy(v).tolist() == [N2]
assert F.asnumpy(e).tolist() == [0]
assert g.in_degree(0, 'follows') == 0
assert g.in_degree(1, 'follows') == 1
assert g.in_degrees(0, 'follows') == 0
assert g.in_degrees(1, 'follows') == 1
assert F.asnumpy(g.in_degrees([0, 1], 'follows')).tolist() == [0, 1]
assert g.in_degree(0, 'plays') == 0
assert g.in_degree(N2, 'plays') == 1
assert g.in_degrees(0, 'plays') == 0
assert g.in_degrees(N2, 'plays') == 1
assert F.asnumpy(g.in_degrees([0, N2], 'plays')).tolist() == [0, 1]
assert g.out_degree(0, 'follows') == 1
assert g.out_degree(1, 'follows') == 0
assert g.out_degrees(0, 'follows') == 1
assert g.out_degrees(1, 'follows') == 0
assert F.asnumpy(g.out_degrees([0, 1], 'follows')).tolist() == [1, 0]
assert g.out_degree(0, 'plays') == 1
assert g.out_degree(N2, 'plays') == 0
assert g.out_degrees(0, 'plays') == 1
assert g.out_degrees(N2, 'plays') == 0
assert F.asnumpy(g.out_degrees([0, N2], 'plays')).tolist() == [1, 0]
def test_edge_ids():
......@@ -380,26 +440,20 @@ def test_edge_ids():
('user', 'follows', 'user'): [(0, 1)],
('user', 'plays', 'game'): [(0, N2)]},
{'user': N1, 'game': N1})
with pytest.raises(AssertionError):
eids = g.edge_ids(0, 0, etype='follows')
with pytest.raises(AssertionError):
eid = g.edge_id(0, 0, etype='follows')
with pytest.raises(DGLError):
eid = g.edge_ids(0, 0, etype='follows')
g2 = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (0, 1)],
('user', 'plays', 'game'): [(0, N2)]},
{'user': N1, 'game': N1})
with pytest.raises(AssertionError):
eids = g2.edge_ids(0, 1, etype='follows')
with pytest.raises(AssertionError):
eid = g2.edge_id(0, 1, etype='follows')
eid = g2.edge_ids(0, 1, etype='follows')
assert eid == 0
@parametrize_dtype
def test_adj(index_dtype):
g = create_test_heterograph(index_dtype)
def test_adj(idtype):
g = create_test_heterograph(idtype)
adj = F.sparse_to_numpy(g.adj(etype='follows'))
assert np.allclose(
adj,
......@@ -454,8 +508,8 @@ def test_adj(index_dtype):
[0., 1., 0.]]))
@parametrize_dtype
def test_inc(index_dtype):
g = create_test_heterograph(index_dtype)
def test_inc(idtype):
g = create_test_heterograph(idtype)
#follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
adj = F.sparse_to_numpy(g['follows'].inc('in'))
assert np.allclose(
......@@ -494,9 +548,9 @@ def test_inc(index_dtype):
[0., 1.]]))
@parametrize_dtype
def test_view(index_dtype):
def test_view(idtype):
# test single node type
g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows', idtype=idtype, device=F.ctx())
f1 = F.randn((3, 6))
g.ndata['h'] = f1
f2 = g.nodes['user'].data['h']
......@@ -521,13 +575,13 @@ def test_view(index_dtype):
assert fail
# test data view
g = create_test_heterograph(index_dtype)
g = create_test_heterograph(idtype)
f1 = F.randn((3, 6))
g.nodes['user'].data['h'] = f1 # ok
f2 = g.nodes['user'].data['h']
assert F.array_equal(f1, f2)
assert F.array_equal(F.tensor(g.nodes('user')), F.arange(0, 3))
assert F.array_equal(g.nodes('user'), F.arange(0, 3, idtype))
g.nodes['user'].data.pop('h')
# multi type ndata
......@@ -560,7 +614,7 @@ def test_view(index_dtype):
f5 = g.edges['follows'].data['h']
assert F.array_equal(f3, f4)
assert F.array_equal(f3, f5)
assert F.array_equal(F.tensor(g.edges(etype='follows', form='eid')), F.arange(0, 2))
assert F.array_equal(g.edges(etype='follows', form='eid'), F.arange(0, 2, idtype))
g.edges['follows'].data.pop('h')
f3 = F.randn((2, 4))
......@@ -588,7 +642,7 @@ def test_view(index_dtype):
g.srcnodes['user'].data['h'] = f1 # ok
f2 = g.srcnodes['user'].data['h']
assert F.array_equal(f1, f2)
assert F.array_equal(F.tensor(g.srcnodes('user')), F.arange(0, 3))
assert F.array_equal(g.srcnodes('user'), F.arange(0, 3, idtype))
g.srcnodes['user'].data.pop('h')
# multi type ndata
......@@ -618,7 +672,7 @@ def test_view(index_dtype):
g.dstnodes['user'].data['h'] = f1 # ok
f2 = g.dstnodes['user'].data['h']
assert F.array_equal(f1, f2)
assert F.array_equal(F.tensor(g.dstnodes('user')), F.arange(0, 3))
assert F.array_equal(g.dstnodes('user'), F.arange(0, 3, idtype))
g.dstnodes['user'].data.pop('h')
# multi type ndata
......@@ -644,9 +698,9 @@ def test_view(index_dtype):
g.dstdata.pop('h')
@parametrize_dtype
def test_view1(index_dtype):
def test_view1(idtype):
# test relation view
HG = create_test_heterograph(index_dtype)
HG = create_test_heterograph(idtype)
ntypes = ['user', 'game', 'developer']
canonical_etypes = [
('user', 'follows', 'user'),
......@@ -661,12 +715,12 @@ def test_view1(index_dtype):
g = HG[etype]
srcs, dsts = edges[etype]
for src, dst in zip(srcs, dsts):
assert g.has_edge_between(src, dst)
assert g.has_edges_between(src, dst)
assert F.asnumpy(g.has_edges_between(srcs, dsts)).all()
srcs, dsts = negative_edges[etype]
for src, dst in zip(srcs, dsts):
assert not g.has_edge_between(src, dst)
assert not g.has_edges_between(src, dst)
assert not F.asnumpy(g.has_edges_between(srcs, dsts)).any()
srcs, dsts = edges[etype]
......@@ -678,7 +732,7 @@ def test_view1(index_dtype):
u, v = g.in_edges([0])
assert F.asnumpy(v).tolist() == [0] * len(pred)
assert set(F.asnumpy(u).tolist()) == set(pred)
assert g.in_degree(0) == len(pred)
assert g.in_degrees(0) == len(pred)
# successors & out_edges & out_degree
succ = [d for s, d in zip(srcs, dsts) if s == 0]
......@@ -686,17 +740,18 @@ def test_view1(index_dtype):
u, v = g.out_edges([0])
assert F.asnumpy(u).tolist() == [0] * len(succ)
assert set(F.asnumpy(v).tolist()) == set(succ)
assert g.out_degree(0) == len(succ)
assert g.out_degrees(0) == len(succ)
# edge_id & edge_ids
for i, (src, dst) in enumerate(zip(srcs, dsts)):
assert g.edge_id(src, dst) == i
assert F.asnumpy(g.edge_id(src, dst, return_array=True)).tolist() == [i]
assert g.edge_ids(src, dst, etype=etype) == i
_, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
assert eid == i
assert F.asnumpy(g.edge_ids(srcs, dsts)).tolist() == list(range(n_edges))
u, v, e = g.edge_ids(srcs, dsts, return_uv=True)
assert F.asnumpy(u).tolist() == srcs
assert F.asnumpy(v).tolist() == dsts
assert F.asnumpy(e).tolist() == list(range(n_edges))
u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
assert u[e].tolist() == srcs
assert v[e].tolist() == dsts
# find_edges
u, v = g.find_edges(list(range(n_edges)))
......@@ -763,13 +818,13 @@ def test_view1(index_dtype):
g.ndata['h'] = f1 # ok
f2 = HG.nodes['user'].data['h']
assert F.array_equal(f1, f2)
assert F.array_equal(F.tensor(g.nodes()), F.arange(0, 3))
assert F.array_equal(g.nodes(), F.arange(0, 3, g.idtype))
f3 = F.randn((2, 4))
g.edata['h'] = f3
f4 = HG.edges['follows'].data['h']
assert F.array_equal(f3, f4)
assert F.array_equal(F.tensor(g.edges(form='eid')), F.arange(0, 2))
assert F.array_equal(g.edges(form='eid'), F.arange(0, 2, g.idtype))
# multiple types
ndata = HG.ndata['h']
......@@ -781,7 +836,7 @@ def test_view1(index_dtype):
assert F.array_equal(edata[('user', 'follows', 'user')], f4)
@parametrize_dtype
def test_flatten(index_dtype):
def test_flatten(idtype):
def check_mapping(g, fg):
if len(fg.ntypes) == 1:
SRC = DST = fg.ntypes[0]
......@@ -805,7 +860,7 @@ def test_flatten(index_dtype):
assert g.canonical_etypes[etype][2] == g.ntypes[tid]
# check for wildcard slices
g = create_test_heterograph(index_dtype)
g = create_test_heterograph(idtype)
g.nodes['user'].data['h'] = F.ones((3, 5))
g.nodes['game'].data['i'] = F.ones((2, 5))
g.edges['plays'].data['e'] = F.ones((4, 4))
......@@ -816,6 +871,8 @@ def test_flatten(index_dtype):
assert len(fg.ntypes) == 2
assert fg.ntypes == ['user', 'game']
assert fg.etypes == ['plays+wishes']
assert fg.idtype == g.idtype
assert fg.device == g.device
assert F.array_equal(fg.nodes['user'].data['h'], F.ones((3, 5)))
assert F.array_equal(fg.nodes['game'].data['i'], F.ones((2, 5)))
......@@ -829,6 +886,8 @@ def test_flatten(index_dtype):
check_mapping(g, fg)
fg = g['user', :, 'user']
assert fg.idtype == g.idtype
assert fg.device == g.device
# NOTE(gq): The node/edge types from the parent graph is returned if there is only one
# node/edge type. This differs from the behavior above.
assert fg.ntypes == ['user']
......@@ -839,6 +898,8 @@ def test_flatten(index_dtype):
assert F.array_equal(v1, v2)
fg = g['developer', :, 'game']
assert fg.idtype == g.idtype
assert fg.device == g.device
assert fg.ntypes == ['developer', 'game']
assert fg.etypes == ['develops']
u1, v1 = g.edges(etype='develops', order='eid')
......@@ -847,13 +908,15 @@ def test_flatten(index_dtype):
assert F.array_equal(v1, v2)
fg = g[:, :, :]
assert fg.idtype == g.idtype
assert fg.device == g.device
assert fg.ntypes == ['developer+user', 'game+user']
assert fg.etypes == ['develops+follows+plays+wishes']
check_mapping(g, fg)
# Test another heterograph
g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user', 'follows', index_dtype=index_dtype)
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', index_dtype=index_dtype)
g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user', 'follows', idtype=idtype, device=F.ctx())
g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', idtype=idtype, device=F.ctx())
g_x.nodes['user'].data['h'] = F.randn((4, 3))
g_x.edges['follows'].data['w'] = F.randn((3, 2))
g_y.nodes['user'].data['hh'] = F.randn((4, 5))
......@@ -866,50 +929,84 @@ def test_flatten(index_dtype):
assert F.array_equal(g.edges['knows'].data['ww'], g_y.edata['ww'])
fg = g['user', :, 'user']
assert fg.idtype == g.idtype
assert fg.device == g.device
assert fg.ntypes == ['user']
assert fg.etypes == ['follows+knows']
check_mapping(g, fg)
fg = g['user', :, :]
assert fg.idtype == g.idtype
assert fg.device == g.device
assert fg.ntypes == ['user']
assert fg.etypes == ['follows+knows']
check_mapping(g, fg)
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
@parametrize_dtype
def test_to_device(index_dtype):
g = create_test_heterograph(index_dtype)
g.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cpu())
g.nodes['game'].data['i'] = F.copy_to(F.ones((2, 5)), F.cpu())
g.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cpu())
def test_to_device(idtype):
# TODO: rewrite this test case to accept different graphs so we
# can test reverse graph and batched graph
g = create_test_heterograph(idtype)
g.nodes['user'].data['h'] = F.ones((3, 5))
g.nodes['game'].data['i'] = F.ones((2, 5))
g.edges['plays'].data['e'] = F.ones((4, 4))
assert g.device == F.ctx()
g = g.to(F.cpu())
assert g.device == F.cpu()
assert F.context(g.nodes['user'].data['h']) == F.cpu()
assert F.context(g.nodes['game'].data['i']) == F.cpu()
assert F.context(g.edges['plays'].data['e']) == F.cpu()
for ntype in g.ntypes:
assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
for etype in g.canonical_etypes:
assert F.context(g.batch_num_edges(etype)) == F.cpu()
if F.is_cuda_available():
g1 = g.to(F.cuda())
assert g1 is not None
assert g1.device == F.cuda()
assert F.context(g1.nodes['user'].data['h']) == F.cuda()
assert F.context(g1.nodes['game'].data['i']) == F.cuda()
assert F.context(g1.edges['plays'].data['e']) == F.cuda()
for ntype in g1.ntypes:
assert F.context(g1.batch_num_nodes(ntype)) == F.cuda()
for etype in g1.canonical_etypes:
assert F.context(g1.batch_num_edges(etype)) == F.cuda()
assert F.context(g.nodes['user'].data['h']) == F.cpu()
assert F.context(g.nodes['game'].data['i']) == F.cpu()
assert F.context(g.edges['plays'].data['e']) == F.cpu()
for ntype in g.ntypes:
assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
for etype in g.canonical_etypes:
assert F.context(g.batch_num_edges(etype)) == F.cpu()
with pytest.raises(DGLError):
g1.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cpu())
with pytest.raises(DGLError):
g1.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cpu())
# set feature after g.to
g = create_test_heterograph(index_dtype)
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
@parametrize_dtype
@pytest.mark.parametrize('g', get_cases(['block']))
def test_to_device2(g, idtype):
g = g.astype(idtype)
g = g.to(F.cpu())
assert g.device == F.cpu()
if F.is_cuda_available():
g1 = g.to(F.cuda())
assert g1 is not None
g1.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cuda())
g1.nodes['game'].data['i'] = F.copy_to(F.ones((2, 5)), F.cuda())
g1.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cuda())
assert g1.device == F.cuda()
assert g1.ntypes == g.ntypes
assert g1.etypes == g.etypes
assert g1.canonical_etypes == g.canonical_etypes
@parametrize_dtype
def test_convert_bound(index_dtype):
def test_convert_bound(idtype):
def _test_bipartite_bound(data, card):
try:
dgl.bipartite(data, num_nodes=card, index_dtype=index_dtype)
except dgl.DGLError:
return
assert False, 'bipartite bound test with wrong uid failed'
with pytest.raises(DGLError):
dgl.bipartite(data, num_nodes=card, idtype=idtype, device=F.ctx())
def _test_graph_bound(data, card):
try:
dgl.graph(data, num_nodes=card, index_dtype=index_dtype)
except dgl.DGLError:
return
assert False, 'graph bound test with wrong uid failed'
with pytest.raises(DGLError):
dgl.graph(data, num_nodes=card, idtype=idtype, device=F.ctx())
_test_bipartite_bound(([1,2],[1,2]),(2,3))
_test_bipartite_bound(([0,1],[1,4]),(2,3))
......@@ -918,8 +1015,8 @@ def test_convert_bound(index_dtype):
@parametrize_dtype
def test_convert(index_dtype):
hg = create_test_heterograph(index_dtype)
def test_convert(idtype):
hg = create_test_heterograph(idtype)
hs = []
for ntype in hg.ntypes:
h = F.randn((hg.number_of_nodes(ntype), 5))
......@@ -934,7 +1031,8 @@ def test_convert(index_dtype):
hg.edges['plays'].data['x'] = F.randn((4, 3))
g = dgl.to_homo(hg)
assert g._idtype_str == index_dtype
assert g.idtype == idtype
assert g.device == hg.device
assert F.array_equal(F.cat(hs, dim=0), g.ndata['h'])
assert 'x' not in g.ndata
assert F.array_equal(F.cat(ws, dim=0), g.edata['w'])
......@@ -963,6 +1061,8 @@ def test_convert(index_dtype):
hg2 = dgl.to_hetero(
g, hg.ntypes, hg.etypes,
ntype_field=dgl.NTYPE, etype_field=dgl.ETYPE, metagraph=_mg)
assert hg2.idtype == hg.idtype
assert hg2.device == hg.device
assert set(hg.ntypes) == set(hg2.ntypes)
assert set(hg.canonical_etypes) == set(hg2.canonical_etypes)
for ntype in hg.ntypes:
......@@ -976,11 +1076,12 @@ def test_convert(index_dtype):
assert F.array_equal(hg.edges[canonical_etype].data['w'], hg2.edges[canonical_etype].data['w'])
# hetero_from_homo test case 2
g = dgl.graph([(0, 2), (1, 2), (2, 3), (0, 3)], index_dtype=index_dtype)
g = dgl.graph([(0, 2), (1, 2), (2, 3), (0, 3)], idtype=idtype, device=F.ctx())
g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2])
g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2])
hg = dgl.to_hetero(g, ['l0', 'l1', 'l2'], ['e0', 'e1', 'e2'])
assert hg._idtype_str == index_dtype
assert hg.idtype == idtype
assert hg.device == g.device
assert set(hg.canonical_etypes) == set(
[('l0', 'e0', 'l1'), ('l1', 'e1', 'l2'), ('l0', 'e2', 'l2')])
assert hg.number_of_nodes('l0') == 2
......@@ -994,12 +1095,13 @@ def test_convert(index_dtype):
mg = nx.MultiDiGraph([
('user', 'movie', 'watches'),
('user', 'TV', 'watches')])
g = dgl.graph([(0, 1), (0, 2)], index_dtype=index_dtype)
g = dgl.graph([(0, 1), (0, 2)], idtype=idtype, device=F.ctx())
g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2])
g.edata[dgl.ETYPE] = F.tensor([0, 0])
for _mg in [None, mg]:
hg = dgl.to_hetero(g, ['user', 'TV', 'movie'], ['watches'], metagraph=_mg)
assert hg._idtype_str == index_dtype
assert hg.idtype == g.idtype
assert hg.device == g.device
assert set(hg.canonical_etypes) == set(
[('user', 'watches', 'movie'), ('user', 'watches', 'TV')])
assert hg.number_of_nodes('user') == 1
......@@ -1010,32 +1112,34 @@ def test_convert(index_dtype):
assert len(hg.etypes) == 2
# hetero_to_homo test case 2
hg = dgl.bipartite([(0, 0), (1, 1)], num_nodes=(2, 3))
hg = dgl.bipartite([(0, 0), (1, 1)], num_nodes=(2, 3), idtype=idtype, device=F.ctx())
g = dgl.to_homo(hg)
assert hg.idtype == g.idtype
assert hg.device == g.device
assert g.number_of_nodes() == 5
@parametrize_dtype
def test_transform(index_dtype):
g = create_test_heterograph(index_dtype)
def test_metagraph_reachable(idtype):
g = create_test_heterograph(idtype)
x = F.randn((3, 5))
g.nodes['user'].data['h'] = x
new_g = dgl.metapath_reachable_graph(g, ['follows', 'plays'])
assert new_g._idtype_str == index_dtype
assert new_g.idtype == idtype
assert new_g.ntypes == ['user', 'game']
assert new_g.number_of_edges() == 3
assert F.asnumpy(new_g.has_edges_between([0, 0, 1], [0, 1, 1])).all()
new_g = dgl.metapath_reachable_graph(g, ['follows'])
assert new_g._idtype_str == index_dtype
assert new_g.idtype == idtype
assert new_g.ntypes == ['user']
assert new_g.number_of_edges() == 2
assert F.asnumpy(new_g.has_edges_between([0, 1], [1, 2])).all()
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support bool tensor")
@parametrize_dtype
def test_subgraph_mask(index_dtype):
g = create_test_heterograph(index_dtype)
def test_subgraph_mask(idtype):
g = create_test_heterograph(idtype)
g_graph = g['follows']
g_bipartite = g['plays']
......@@ -1045,36 +1149,39 @@ def test_subgraph_mask(index_dtype):
g.edges['follows'].data['h'] = y
def _check_subgraph(g, sg):
assert sg.idtype == g.idtype
assert sg.device == g.device
assert sg.ntypes == g.ntypes
assert sg.etypes == g.etypes
assert sg.canonical_etypes == g.canonical_etypes
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
F.tensor([1, 2], F.int64))
F.tensor([1, 2], idtype))
assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
F.tensor([0], F.int64))
F.tensor([0], idtype))
assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
F.tensor([1], F.int64))
F.tensor([1], idtype))
assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
F.tensor([1], F.int64))
F.tensor([1], idtype))
assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
F.tensor([1], F.int64))
F.tensor([1], idtype))
assert sg.number_of_nodes('developer') == 0
assert sg.number_of_edges('develops') == 0
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'][1:3])
assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'][1:2])
# backend boo input tensor
sg1 = g.subgraph({'user': F.tensor([False, True, True], dtype=F.data_type_dict['bool']),
'game': F.tensor([True, False, False, False], dtype=F.data_type_dict['bool'])})
sg1 = g.subgraph({'user': F.tensor([False, True, True], dtype=F.bool),
'game': F.tensor([True, False, False, False], dtype=F.bool)})
_check_subgraph(g, sg1)
sg2 = g.edge_subgraph({'follows': F.tensor([False, True], dtype=F.data_type_dict['bool']),
'plays': F.tensor([False, True, False, False], dtype=F.data_type_dict['bool']),
'wishes': F.tensor([False, True], dtype=F.data_type_dict['bool'])})
_check_subgraph(g, sg2)
if F._default_context_str != 'gpu':
# TODO(minjie): enable this later
sg2 = g.edge_subgraph({'follows': F.tensor([False, True], dtype=F.bool),
'plays': F.tensor([False, True, False, False], dtype=F.bool),
'wishes': F.tensor([False, True], dtype=F.bool)})
_check_subgraph(g, sg2)
@parametrize_dtype
def test_subgraph(index_dtype):
g = create_test_heterograph(index_dtype)
def test_subgraph(idtype):
g = create_test_heterograph(idtype)
g_graph = g['follows']
g_bipartite = g['plays']
......@@ -1084,19 +1191,21 @@ def test_subgraph(index_dtype):
g.edges['follows'].data['h'] = y
def _check_subgraph(g, sg):
assert sg.idtype == g.idtype
assert sg.device == g.device
assert sg.ntypes == g.ntypes
assert sg.etypes == g.etypes
assert sg.canonical_etypes == g.canonical_etypes
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
F.tensor([1, 2], F.int64))
F.tensor([1, 2], g.idtype))
assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
F.tensor([0], F.int64))
F.tensor([0], g.idtype))
assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
F.tensor([1], F.int64))
F.tensor([1], g.idtype))
assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
F.tensor([1], F.int64))
F.tensor([1], g.idtype))
assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
F.tensor([1], F.int64))
F.tensor([1], g.idtype))
assert sg.number_of_nodes('developer') == 0
assert sg.number_of_edges('develops') == 0
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'][1:3])
......@@ -1104,41 +1213,49 @@ def test_subgraph(index_dtype):
sg1 = g.subgraph({'user': [1, 2], 'game': [0]})
_check_subgraph(g, sg1)
sg2 = g.edge_subgraph({'follows': [1], 'plays': [1], 'wishes': [1]})
_check_subgraph(g, sg2)
if F._default_context_str != 'gpu':
# TODO(minjie): enable this later
sg2 = g.edge_subgraph({'follows': [1], 'plays': [1], 'wishes': [1]})
_check_subgraph(g, sg2)
# backend tensor input
sg1 = g.subgraph({'user': F.tensor([1, 2], dtype=F.data_type_dict[index_dtype]),
'game': F.tensor([0], dtype=F.data_type_dict[index_dtype])})
sg1 = g.subgraph({'user': F.tensor([1, 2], dtype=idtype),
'game': F.tensor([0], dtype=idtype)})
_check_subgraph(g, sg1)
sg2 = g.edge_subgraph({'follows': F.tensor([1], dtype=F.data_type_dict[index_dtype]),
'plays': F.tensor([1], dtype=F.data_type_dict[index_dtype]),
'wishes': F.tensor([1], dtype=F.data_type_dict[index_dtype])})
_check_subgraph(g, sg2)
if F._default_context_str != 'gpu':
# TODO(minjie): enable this later
sg2 = g.edge_subgraph({'follows': F.tensor([1], dtype=idtype),
'plays': F.tensor([1], dtype=idtype),
'wishes': F.tensor([1], dtype=idtype)})
_check_subgraph(g, sg2)
# numpy input
sg1 = g.subgraph({'user': np.array([1, 2]),
'game': np.array([0])})
_check_subgraph(g, sg1)
sg2 = g.edge_subgraph({'follows': np.array([1]),
'plays': np.array([1]),
'wishes': np.array([1])})
_check_subgraph(g, sg2)
if F._default_context_str != 'gpu':
# TODO(minjie): enable this later
sg2 = g.edge_subgraph({'follows': np.array([1]),
'plays': np.array([1]),
'wishes': np.array([1])})
_check_subgraph(g, sg2)
def _check_subgraph_single_ntype(g, sg, preserve_nodes=False):
assert sg.idtype == g.idtype
assert sg.device == g.device
assert sg.ntypes == g.ntypes
assert sg.etypes == g.etypes
assert sg.canonical_etypes == g.canonical_etypes
if not preserve_nodes:
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
F.tensor([1, 2], F.int64))
F.tensor([1, 2], g.idtype))
else:
for ntype in sg.ntypes:
assert g.number_of_nodes(ntype) == sg.number_of_nodes(ntype)
assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
F.tensor([1], F.int64))
F.tensor([1], g.idtype))
if not preserve_nodes:
assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'][1:3])
......@@ -1151,28 +1268,32 @@ def test_subgraph(index_dtype):
if not preserve_nodes:
assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
F.tensor([0, 1], F.int64))
F.tensor([0, 1], g.idtype))
assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
F.tensor([0], F.int64))
F.tensor([0], g.idtype))
else:
for ntype in sg.ntypes:
assert g.number_of_nodes(ntype) == sg.number_of_nodes(ntype)
assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
F.tensor([0, 1], F.int64))
F.tensor([0, 1], g.idtype))
sg1_graph = g_graph.subgraph([1, 2])
_check_subgraph_single_ntype(g_graph, sg1_graph)
sg1_graph = g_graph.edge_subgraph([1])
_check_subgraph_single_ntype(g_graph, sg1_graph)
sg1_graph = g_graph.edge_subgraph([1], preserve_nodes=True)
_check_subgraph_single_ntype(g_graph, sg1_graph, True)
sg2_bipartite = g_bipartite.edge_subgraph([0, 1])
_check_subgraph_single_etype(g_bipartite, sg2_bipartite)
sg2_bipartite = g_bipartite.edge_subgraph([0, 1], preserve_nodes=True)
_check_subgraph_single_etype(g_bipartite, sg2_bipartite, True)
if F._default_context_str != 'gpu':
# TODO(minjie): enable this later
sg1_graph = g_graph.edge_subgraph([1])
_check_subgraph_single_ntype(g_graph, sg1_graph)
sg1_graph = g_graph.edge_subgraph([1], preserve_nodes=True)
_check_subgraph_single_ntype(g_graph, sg1_graph, True)
sg2_bipartite = g_bipartite.edge_subgraph([0, 1])
_check_subgraph_single_etype(g_bipartite, sg2_bipartite)
sg2_bipartite = g_bipartite.edge_subgraph([0, 1], preserve_nodes=True)
_check_subgraph_single_etype(g_bipartite, sg2_bipartite, True)
def _check_typed_subgraph1(g, sg):
assert g.idtype == sg.idtype
assert g.device == sg.device
assert set(sg.ntypes) == {'user', 'game'}
assert set(sg.etypes) == {'follows', 'plays', 'wishes'}
for ntype in sg.ntypes:
......@@ -1208,25 +1329,27 @@ def test_subgraph(index_dtype):
_check_typed_subgraph1(g, sg5)
# Test for restricted format
for fmt in ['csr', 'csc', 'coo']:
g = dgl.graph([(0, 1), (1, 2)], restrict_format=fmt)
sg = g.subgraph({g.ntypes[0]: [1, 0]})
nids = F.asnumpy(sg.ndata[dgl.NID])
assert np.array_equal(nids, np.array([1, 0]))
src, dst = sg.all_edges(order='eid')
src = F.asnumpy(src)
dst = F.asnumpy(dst)
assert np.array_equal(src, np.array([1]))
assert np.array_equal(dst, np.array([0]))
if F._default_context_str != 'gpu':
# TODO(minjie): enable this later
for fmt in ['csr', 'csc', 'coo']:
g = dgl.graph([(0, 1), (1, 2)]).formats(fmt)
sg = g.subgraph({g.ntypes[0]: [1, 0]})
nids = F.asnumpy(sg.ndata[dgl.NID])
assert np.array_equal(nids, np.array([1, 0]))
src, dst = sg.edges(order='eid')
src = F.asnumpy(src)
dst = F.asnumpy(dst)
assert np.array_equal(src, np.array([1]))
assert np.array_equal(dst, np.array([0]))
@parametrize_dtype
def test_apply(index_dtype):
def test_apply(idtype):
def node_udf(nodes):
return {'h': nodes.data['h'] * 2}
def edge_udf(edges):
return {'h': edges.data['h'] * 2 + edges.src['h']}
g = create_test_heterograph(index_dtype)
g = create_test_heterograph(idtype)
g.nodes['user'].data['h'] = F.ones((3, 5))
g.apply_nodes(node_udf, ntype='user')
assert F.array_equal(g.nodes['user'].data['h'], F.ones((3, 5)) * 2)
......@@ -1244,122 +1367,21 @@ def test_apply(index_dtype):
# test fail case
# fail due to multiple types
fail = False
try:
with pytest.raises(DGLError):
g.apply_nodes(node_udf)
except dgl.DGLError:
fail = True
assert fail
fail = False
try:
with pytest.raises(DGLError):
g.apply_edges(edge_udf)
except dgl.DGLError:
fail = True
assert fail
@parametrize_dtype
def test_level1(index_dtype):
def test_level2(idtype):
#edges = {
# 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]),
#}
g = create_test_heterograph(index_dtype)
def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)}
def rfunc2(nodes):
return {'y': F.max(nodes.mailbox['m'], 1)}
def mfunc(edges):
return {'m': edges.src['h']}
def afunc(nodes):
return {'y' : nodes.data['y'] + 1}
g.nodes['user'].data['h'] = F.ones((3, 2))
g.send([2, 3], mfunc, etype='plays')
g.recv([0, 1], rfunc, etype='plays')
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
g.nodes['game'].data.pop('y')
# only one type
play_g = g['plays']
play_g.send([2, 3], mfunc)
play_g.recv([0, 1], rfunc)
y = g.nodes['game'].data['y']
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
# TODO(minjie): following codes will fail because messages are
# not shared with the base graph. However, since send and recv
# are rarely used, no fix at the moment.
# g['plays'].send([2, 3], mfunc)
# g['plays'].recv([0, 1], mfunc)
# test fail case
# fail due to multiple types
fail = False
try:
g.send([2, 3], mfunc)
except dgl.DGLError:
fail = True
assert fail
fail = False
try:
g.recv([0, 1], rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi recv
g.send(g.edges(etype='plays'), mfunc, etype='plays')
g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
g.multi_recv([0, 1], {'plays' : rfunc, ('user', 'wishes', 'game'): rfunc2}, 'sum')
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[3., 3.], [3., 3.]]))
# test multi recv with apply function
g.send(g.edges(etype='plays'), mfunc, etype='plays')
g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
g.multi_recv([0, 1], {'plays' : (rfunc, afunc), ('user', 'wishes', 'game'): rfunc2}, 'sum', afunc)
assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[5., 5.], [5., 5.]]))
# test cross reducer
g.nodes['user'].data['h'] = F.randn((3, 2))
for cred in ['sum', 'max', 'min', 'mean']:
g.send(g.edges(etype='plays'), mfunc, etype='plays')
g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
g.multi_recv([0, 1], {'plays' : (rfunc, afunc), 'wishes': rfunc2}, cred, afunc)
y = g.nodes['game'].data['y']
g1 = g['plays']
g2 = g['wishes']
g1.send(g1.edges(), mfunc)
g1.recv(g1.nodes('game'), rfunc, afunc)
y1 = g.nodes['game'].data['y']
g2.send(g2.edges(), mfunc)
g2.recv(g2.nodes('game'), rfunc2)
y2 = g.nodes['game'].data['y']
yy = get_redfn(cred)(F.stack([y1, y2], 0), 0)
yy = yy + 1 # final afunc
assert F.array_equal(y, yy)
# test fail case
# fail because cannot infer ntype
fail = False
try:
g.multi_recv([0, 1], {'plays' : rfunc, 'follows': rfunc2}, 'sum')
except dgl.DGLError:
fail = True
assert fail
@parametrize_dtype
def test_level2(index_dtype):
#edges = {
# 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]),
#}
g = create_test_heterograph(index_dtype)
g = create_test_heterograph(idtype)
def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)}
def rfunc2(nodes):
......@@ -1385,12 +1407,8 @@ def test_level2(index_dtype):
# test fail case
# fail due to multiple types
fail = False
try:
with pytest.raises(DGLError):
g.send_and_recv([2, 3], mfunc, rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi
g.multi_send_and_recv(
......@@ -1424,15 +1442,11 @@ def test_level2(index_dtype):
# test fail case
# fail because cannot infer ntype
fail = False
try:
with pytest.raises(DGLError):
g.multi_send_and_recv(
{'plays' : (g.edges(etype='plays'), mfunc, rfunc),
'follows': (g.edges(etype='follows'), mfunc, rfunc2)},
'sum')
except dgl.DGLError:
fail = True
assert fail
g.nodes['game'].data.clear()
......@@ -1451,12 +1465,8 @@ def test_level2(index_dtype):
assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
# test fail case
fail = False
try:
with pytest.raises(DGLError):
g.pull(1, mfunc, rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi
g.multi_pull(
......@@ -1494,16 +1504,12 @@ def test_level2(index_dtype):
# test fail case
# fail because cannot infer ntype
fail = False
try:
with pytest.raises(DGLError):
g.multi_pull(
1,
{'plays' : (mfunc, rfunc),
'follows': (mfunc, rfunc2)},
'sum')
except dgl.DGLError:
fail = True
assert fail
g.nodes['game'].data.clear()
......@@ -1523,12 +1529,8 @@ def test_level2(index_dtype):
# test fail case
# fail due to multiple types
fail = False
try:
with pytest.raises(DGLError):
g.update_all(mfunc, rfunc)
except dgl.DGLError:
fail = True
assert fail
# test multi
g.multi_update_all(
......@@ -1568,27 +1570,23 @@ def test_level2(index_dtype):
# test fail case
# fail because cannot infer ntype
fail = False
try:
with pytest.raises(DGLError):
g.update_all(
{'plays' : (mfunc, rfunc),
'follows': (mfunc, rfunc2)},
'sum')
except dgl.DGLError:
fail = True
assert fail
g.nodes['game'].data.clear()
@parametrize_dtype
def test_updates(index_dtype):
def test_updates(idtype):
def msg_func(edges):
return {'m': edges.src['h']}
def reduce_func(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)}
def apply_func(nodes):
return {'y': nodes.data['y'] * 2}
g = create_test_heterograph(index_dtype)
g = create_test_heterograph(idtype)
x = F.randn((3, 5))
g.nodes['user'].data['h'] = x
......@@ -1609,14 +1607,6 @@ def test_updates(index_dtype):
assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
del g.nodes['game'].data['y']
plays_g = g['user', 'plays', 'game']
plays_g.send(([0, 1, 2], [0, 1, 1]), msg)
plays_g.recv([0, 1], red, apply)
y = g.nodes['game'].data['y']
assert F.array_equal(y[0], x[0] * multiplier)
assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
del g.nodes['game'].data['y']
# pulls from destination (game) node 0
g['user', 'plays', 'game'].pull(0, msg, red, apply)
y = g.nodes['game'].data['y']
......@@ -1631,8 +1621,8 @@ def test_updates(index_dtype):
@parametrize_dtype
def test_backward(index_dtype):
g = create_test_heterograph(index_dtype)
def test_backward(idtype):
g = create_test_heterograph(idtype)
x = F.randn((3, 5))
F.attach_grad(x)
g.nodes['user'].data['h'] = x
......@@ -1650,7 +1640,7 @@ def test_backward(index_dtype):
@parametrize_dtype
def test_empty_heterograph(index_dtype):
def test_empty_heterograph(idtype):
def assert_empty(g):
assert g.number_of_nodes('user') == 0
assert g.number_of_edges('plays') == 0
......@@ -1665,15 +1655,17 @@ def test_empty_heterograph(index_dtype):
# empty networkx graph
assert_empty(dgl.heterograph({('user', 'plays', 'game'): nx.DiGraph()}))
g = dgl.heterograph({('user', 'follows', 'user'): []}, index_dtype=index_dtype)
assert g._idtype_str == index_dtype
g = dgl.heterograph({('user', 'follows', 'user'): []}, idtype=idtype, device=F.ctx())
assert g.idtype == idtype
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 0
assert g.number_of_edges('follows') == 0
# empty relation graph with others
g = dgl.heterograph({('user', 'plays', 'game'): [], ('developer', 'develops', 'game'): [
(0, 0), (1, 1)]}, index_dtype=index_dtype)
assert g._idtype_str == index_dtype
(0, 0), (1, 1)]}, idtype=idtype, device=F.ctx())
assert g.idtype == idtype
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 0
assert g.number_of_edges('plays') == 0
assert g.number_of_nodes('game') == 2
......@@ -1719,8 +1711,6 @@ def test_types_in_function():
g.apply_edges(mfunc1)
g.update_all(mfunc1, rfunc1)
g.send_and_recv([0, 1], mfunc1, rfunc1)
g.send([0, 1], mfunc1)
g.recv([1, 2], rfunc1)
g.push([0], mfunc1, rfunc1)
g.pull([1], mfunc1, rfunc1)
g.filter_nodes(filter_nodes1)
......@@ -1731,22 +1721,20 @@ def test_types_in_function():
g.apply_edges(mfunc2)
g.update_all(mfunc2, rfunc2)
g.send_and_recv([0, 1], mfunc2, rfunc2)
g.send([0, 1], mfunc2)
g.recv([1, 2], rfunc2)
g.push([0], mfunc2, rfunc2)
g.pull([1], mfunc2, rfunc2)
g.filter_nodes(filter_nodes2, ntype='game')
g.filter_edges(filter_edges2)
@parametrize_dtype
def test_stack_reduce(index_dtype):
def test_stack_reduce(idtype):
#edges = {
# 'follows': ([0, 1], [1, 2]),
# 'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
# 'wishes': ([0, 2], [1, 0]),
# 'develops': ([0, 1], [0, 1]),
#}
g = create_test_heterograph(index_dtype)
g = create_test_heterograph(idtype)
g.nodes['user'].data['h'] = F.randn((3, 200))
def rfunc(nodes):
return {'y': F.sum(nodes.mailbox['m'], 1)}
......@@ -1766,24 +1754,24 @@ def test_stack_reduce(index_dtype):
assert g.nodes['game'].data['y'].shape == (g.number_of_nodes('game'), 1, 200)
@parametrize_dtype
def test_isolated_ntype(index_dtype):
def test_isolated_ntype(idtype):
g = dgl.heterograph({
('A', 'AB', 'B'): [(0, 1), (1, 2), (2, 3)]},
num_nodes_dict={'A': 3, 'B': 4, 'C': 4})
num_nodes_dict={'A': 3, 'B': 4, 'C': 4},
idtype=idtype, device=F.ctx())
assert g.number_of_nodes('A') == 3
assert g.number_of_nodes('B') == 4
assert g.number_of_nodes('C') == 4
g = dgl.heterograph({
('A', 'AC', 'C'): [(0, 1), (1, 2), (2, 3)]},
num_nodes_dict={'A': 3, 'B': 4, 'C': 4})
num_nodes_dict={'A': 3, 'B': 4, 'C': 4},
idtype=idtype, device=F.ctx())
assert g.number_of_nodes('A') == 3
assert g.number_of_nodes('B') == 4
assert g.number_of_nodes('C') == 4
G = dgl.DGLGraph()
G.add_nodes(11)
G.add_edges([0, 1, 2], [4, 5, 6])
G = dgl.graph(([0, 1, 2], [4, 5, 6]), num_nodes=11, idtype=idtype, device=F.ctx())
G.ndata[dgl.NTYPE] = F.tensor([0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], dtype=F.int64)
G.edata[dgl.ETYPE] = F.tensor([0, 0, 0], dtype=F.int64)
g = dgl.to_hetero(G, ['A', 'B', 'C'], ['AB'])
......@@ -1793,18 +1781,18 @@ def test_isolated_ntype(index_dtype):
@parametrize_dtype
def test_ismultigraph(index_dtype):
def test_ismultigraph(idtype):
g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5), (2, 5)], 'A',
'AB', 'B', num_nodes=(6, 6), index_dtype=index_dtype)
'AB', 'B', num_nodes=(6, 6), idtype=idtype, device=F.ctx())
assert g1.is_multigraph == False
g2 = dgl.bipartite([(0, 1), (0, 1), (0, 2), (1, 5)], 'A',
'AC', 'C', num_nodes=(6, 6), index_dtype=index_dtype)
'AC', 'C', num_nodes=(6, 6), idtype=idtype, device=F.ctx())
assert g2.is_multigraph == True
g3 = dgl.graph([(0, 1), (1, 2)], 'A', 'AA',
num_nodes=6, index_dtype=index_dtype)
num_nodes=6, idtype=idtype, device=F.ctx())
assert g3.is_multigraph == False
g4 = dgl.graph([(0, 1), (0, 1), (1, 2)], 'A', 'AA',
num_nodes=6, index_dtype=index_dtype)
num_nodes=6, idtype=idtype, device=F.ctx())
assert g4.is_multigraph == True
g = dgl.hetero_from_relations([g1, g3])
assert g.is_multigraph == False
......@@ -1816,8 +1804,8 @@ def test_ismultigraph(index_dtype):
assert g.is_multigraph == True
@parametrize_dtype
def test_bipartite(index_dtype):
g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5)], 'A', 'AB', 'B', index_dtype=index_dtype)
def test_bipartite(idtype):
g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5)], 'A', 'AB', 'B', idtype=idtype, device=F.ctx())
assert g1.is_unibipartite
assert len(g1.ntypes) == 2
assert g1.etypes == ['AB']
......@@ -1840,7 +1828,7 @@ def test_bipartite(index_dtype):
assert F.array_equal(g1.nodes['DST/B'].data['h'], g1.dstdata['h'])
# more complicated bipartite
g2 = dgl.bipartite([(1, 0), (0, 0)], 'A', 'AC', 'C', index_dtype=index_dtype)
g2 = dgl.bipartite([(1, 0), (0, 0)], 'A', 'AC', 'C', idtype=idtype, device=F.ctx())
g3 = dgl.hetero_from_relations([g1, g2])
assert g3.is_unibipartite
assert g3.srctypes == ['A']
......@@ -1857,93 +1845,80 @@ def test_bipartite(index_dtype):
assert F.array_equal(g3.nodes['A'].data['h'], g3.srcdata['h'])
assert F.array_equal(g3.nodes['SRC/A'].data['h'], g3.srcdata['h'])
g4 = dgl.graph([(0, 0), (1, 1)], 'A', 'AA', index_dtype=index_dtype)
g4 = dgl.graph([(0, 0), (1, 1)], 'A', 'AA', idtype=idtype, device=F.ctx())
g5 = dgl.hetero_from_relations([g1, g2, g4])
assert not g5.is_unibipartite
@parametrize_dtype
def test_dtype_cast(index_dtype):
g = dgl.graph([(0, 0), (1, 1), (0, 1), (2, 0)], index_dtype=index_dtype)
assert g._idtype_str == index_dtype
def test_dtype_cast(idtype):
g = dgl.graph([(0, 0), (1, 1), (0, 1), (2, 0)], idtype=idtype, device=F.ctx())
assert g.idtype == idtype
g.ndata["feat"] = F.tensor([3, 4, 5])
g.edata["h"] = F.tensor([3, 4, 5, 6])
if index_dtype == "int32":
if idtype == "int32":
g_cast = g.long()
assert g_cast._idtype_str == 'int64'
assert g_cast.idtype == F.int64
else:
g_cast = g.int()
assert g_cast._idtype_str == 'int32'
assert "feat" in g_cast.ndata
assert "h" in g_cast.edata
assert F.array_equal(g.ndata["feat"], g_cast.ndata["feat"])
assert F.array_equal(g.edata["h"], g_cast.edata["h"])
assert g_cast.idtype == F.int32
test_utils.check_graph_equal(g, g_cast, check_idtype=False)
def test_format():
@parametrize_dtype
def test_format(idtype):
# single relation
g = dgl.graph([(0, 0), (1, 1), (0, 1), (2, 0)], restrict_format='coo')
assert g.restrict_format() == 'coo'
assert g.format_in_use() == ['coo']
g = dgl.graph([(0, 0), (1, 1), (0, 1), (2, 0)], idtype=idtype, device=F.ctx()).formats('coo')
assert g.formats()['created'] == ['coo']
assert len(g.formats()['not created']) == 0
try:
spmat = g.adjacency_matrix(scipy_fmt="csr")
except:
print('test passed, graph with restrict_format coo should not create csr matrix.')
print('test passed, graph with allowed format coo should not create csr matrix.')
else:
assert False, 'cannot create csr when restrict_format is coo'
g1 = g.to_format('any')
assert g1.restrict_format() == 'any'
g1.request_format('coo')
g1.request_format('csr')
g1.request_format('csc')
assert len(g1.format_in_use()) == 3
assert g.restrict_format() == 'coo'
assert g.format_in_use() == ['coo']
assert False, 'cannot create csr when allowed format is coo'
g1 = g.formats(['coo', 'csr', 'csc'])
assert len(g1.formats()['created']) + len(g1.formats()['not created']) == 3
g1.create_format_()
assert len(g1.formats()['created']) == 3
assert g.formats()['created'] == ['coo']
# multiple relation
g = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (1, 2)],
('user', 'plays', 'game'): [(0, 0), (1, 0), (1, 1), (2, 1)],
('developer', 'develops', 'game'): [(0, 0), (1, 1)],
}, restrict_format='csr')
}, idtype=idtype, device=F.ctx()).formats('csr')
user_feat = F.randn((g['follows'].number_of_src_nodes(), 5))
g['follows'].srcdata['h'] = user_feat
for rel_type in ['follows', 'plays', 'develops']:
assert g.restrict_format(rel_type) == 'csr'
assert g.format_in_use(rel_type) == ['csr']
try:
g[rel_type].request_format('coo')
except:
print('test passed, graph with restrict_format csr should not create coo matrix')
else:
assert False, 'cannot create coo when restrict_format is csr'
assert g.formats()['created'] == ['csr']
assert len(g.formats()['not created']) == 0
g1 = g.to_format('csc')
g1 = g.formats('csc')
# test frame
assert F.array_equal(g1['follows'].srcdata['h'], user_feat)
# test each relation graph
for rel_type in ['follows', 'plays', 'develops']:
assert g1.restrict_format(rel_type) == 'csc'
assert g1.format_in_use(rel_type) == ['csc']
assert g.restrict_format(rel_type) == 'csr'
assert g.format_in_use(rel_type) == ['csr']
assert g1.formats()['created'] == ['csc']
assert len(g1.formats()['not created']) == 0
assert g.formats()['created'] == ['csr']
assert len(g.formats()['not created']) == 0
def test_edges_order():
@parametrize_dtype
def test_edges_order(idtype):
# (0, 2), (1, 2), (0, 1), (0, 1), (2, 1)
g = dgl.graph((
np.array([0, 1, 0, 0, 2]),
np.array([2, 2, 1, 1, 1])
))
), idtype=idtype, device=F.ctx())
print(g.formats())
src, dst = g.all_edges(order='srcdst')
assert F.array_equal(F.copy_to(src, F.cpu()),
F.copy_to(F.tensor([0, 0, 0, 1, 2]), F.cpu()))
assert F.array_equal(F.copy_to(dst, F.cpu()),
F.copy_to(F.tensor([1, 1, 2, 2, 1]), F.cpu()))
assert F.array_equal(src, F.tensor([0, 0, 0, 1, 2], dtype=idtype))
assert F.array_equal(dst, F.tensor([1, 1, 2, 2, 1], dtype=idtype))
@parametrize_dtype
def test_reverse(index_dtype):
def test_reverse(idtype):
g = dgl.heterograph({
('user', 'follows', 'user'): ([0, 1, 2, 4, 3 ,1, 3], [1, 2, 3, 2, 0, 0, 1]),
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
gidx = g._graph
r_gidx = gidx.reverse()
......@@ -1951,40 +1926,40 @@ def test_reverse(index_dtype):
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
g_s, g_d, _ = gidx.edges(0)
rg_s, rg_d, _ = r_gidx.edges(0)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
# force to start with 'csr'
gidx = gidx.to_format('csr')
gidx = gidx.to_format('any')
gidx = gidx.formats('csr')
gidx = gidx.formats(['coo', 'csr', 'csc'])
r_gidx = gidx.reverse()
assert gidx.format_in_use(0)[0] == 'csr'
assert r_gidx.format_in_use(0)[0] == 'csc'
assert 'csr' in gidx.formats()['created']
assert 'csc' in r_gidx.formats()['created']
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
g_s, g_d, _ = gidx.edges(0)
rg_s, rg_d, _ = r_gidx.edges(0)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
# force to start with 'csc'
gidx = gidx.to_format('csc')
gidx = gidx.to_format('any')
gidx = gidx.formats('csc')
gidx = gidx.formats(['coo', 'csr', 'csc'])
r_gidx = gidx.reverse()
assert gidx.format_in_use(0)[0] == 'csc'
assert r_gidx.format_in_use(0)[0] == 'csr'
assert 'csc' in gidx.formats()['created']
assert 'csr' in r_gidx.formats()['created']
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
g_s, g_d, _ = gidx.edges(0)
rg_s, rg_d, _ = r_gidx.edges(0)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
g = dgl.heterograph({
('user', 'follows', 'user'): ([0, 1, 2, 4, 3 ,1, 3], [1, 2, 3, 2, 0, 0, 1]),
('user', 'plays', 'game'): ([0, 0, 2, 3, 3, 4, 1], [1, 0, 1, 0, 1, 0, 0]),
('developer', 'develops', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1]),
}, index_dtype=index_dtype)
}, idtype=idtype, device=F.ctx())
gidx = g._graph
r_gidx = gidx.reverse()
......@@ -2003,28 +1978,24 @@ def test_reverse(index_dtype):
assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
g_s, g_d, _ = gidx.edges(0)
rg_s, rg_d, _ = r_gidx.edges(0)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
g_s, g_d, _ = gidx.edges(1)
rg_s, rg_d, _ = r_gidx.edges(1)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
g_s, g_d, _ = gidx.edges(2)
rg_s, rg_d, _ = r_gidx.edges(2)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
# force to start with 'csr'
gidx = gidx.to_format('csr')
gidx = gidx.to_format('any')
gidx = gidx.formats('csr')
gidx = gidx.formats(['coo', 'csr', 'csc'])
r_gidx = gidx.reverse()
# three node types and three edge types
assert gidx.format_in_use(0)[0] == 'csr'
assert r_gidx.format_in_use(0)[0] == 'csc'
assert gidx.format_in_use(1)[0] == 'csr'
assert r_gidx.format_in_use(1)[0] == 'csc'
assert gidx.format_in_use(2)[0] == 'csr'
assert r_gidx.format_in_use(2)[0] == 'csc'
assert 'csr' in gidx.formats()['created']
assert 'csc' in r_gidx.formats()['created']
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
......@@ -2033,28 +2004,24 @@ def test_reverse(index_dtype):
assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
g_s, g_d, _ = gidx.edges(0)
rg_s, rg_d, _ = r_gidx.edges(0)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
g_s, g_d, _ = gidx.edges(1)
rg_s, rg_d, _ = r_gidx.edges(1)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
g_s, g_d, _ = gidx.edges(2)
rg_s, rg_d, _ = r_gidx.edges(2)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
# force to start with 'csc'
gidx = gidx.to_format('csc')
gidx = gidx.to_format('any')
gidx = gidx.formats('csc')
gidx = gidx.formats(['coo', 'csr', 'csc'])
r_gidx = gidx.reverse()
# three node types and three edge types
assert gidx.format_in_use(0)[0] == 'csc'
assert r_gidx.format_in_use(0)[0] == 'csr'
assert gidx.format_in_use(1)[0] == 'csc'
assert r_gidx.format_in_use(1)[0] == 'csr'
assert gidx.format_in_use(2)[0] == 'csc'
assert r_gidx.format_in_use(2)[0] == 'csr'
assert 'csc' in gidx.formats()['created']
assert 'csr' in r_gidx.formats()['created']
assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
......@@ -2063,16 +2030,462 @@ def test_reverse(index_dtype):
assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
g_s, g_d, _ = gidx.edges(0)
rg_s, rg_d, _ = r_gidx.edges(0)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
g_s, g_d, _ = gidx.edges(1)
rg_s, rg_d, _ = r_gidx.edges(1)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
g_s, g_d, _ = gidx.edges(2)
rg_s, rg_d, _ = r_gidx.edges(2)
assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
assert F.array_equal(g_s, rg_d)
assert F.array_equal(g_d, rg_s)
@parametrize_dtype
def test_clone(idtype):
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
new_g = g.clone()
assert g.number_of_nodes() == new_g.number_of_nodes()
assert g.number_of_edges() == new_g.number_of_edges()
assert g.device == new_g.device
assert g.idtype == new_g.idtype
assert F.array_equal(g.ndata['h'], new_g.ndata['h'])
assert F.array_equal(g.edata['h'], new_g.edata['h'])
# data change
new_g.ndata['h'] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())
assert (F.array_equal(g.ndata['h'], new_g.ndata['h']) == False)
g.edata['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
assert (F.array_equal(g.edata['h'], new_g.edata['h']) == False)
# graph structure change
g.add_nodes(1)
assert g.number_of_nodes() != new_g.number_of_nodes()
new_g.add_edges(1, 1)
assert g.number_of_edges() != new_g.number_of_edges()
# zero data graph
g = dgl.graph([], num_nodes=0, idtype=idtype, device=F.ctx())
new_g = g.clone()
assert g.number_of_nodes() == new_g.number_of_nodes()
assert g.number_of_edges() == new_g.number_of_edges()
# heterograph
g = create_test_heterograph4(idtype)
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
new_g = g.clone()
assert g.number_of_nodes('user') == new_g.number_of_nodes('user')
assert g.number_of_nodes('game') == new_g.number_of_nodes('game')
assert g.number_of_nodes('developer') == new_g.number_of_nodes('developer')
assert g.number_of_edges('plays') == new_g.number_of_edges('plays')
assert g.number_of_edges('develops') == new_g.number_of_edges('develops')
assert F.array_equal(g.nodes['user'].data['h'], new_g.nodes['user'].data['h'])
assert F.array_equal(g.nodes['game'].data['h'], new_g.nodes['game'].data['h'])
assert F.array_equal(g.edges['plays'].data['h'], new_g.edges['plays'].data['h'])
assert g.device == new_g.device
assert g.idtype == new_g.idtype
u, v = g.edges(form='uv', order='eid', etype='plays')
nu, nv = new_g.edges(form='uv', order='eid', etype='plays')
assert F.array_equal(u, nu)
assert F.array_equal(v, nv)
# graph structure change
u = F.tensor([0, 4], dtype=idtype)
v = F.tensor([2, 6], dtype=idtype)
g.add_edges(u, v, etype='plays')
u, v = g.edges(form='uv', order='eid', etype='plays')
assert u.shape[0] != nu.shape[0]
assert v.shape[0] != nv.shape[0]
assert g.nodes['user'].data['h'].shape[0] != new_g.nodes['user'].data['h'].shape[0]
assert g.nodes['game'].data['h'].shape[0] != new_g.nodes['game'].data['h'].shape[0]
assert g.edges['plays'].data['h'].shape[0] != new_g.edges['plays'].data['h'].shape[0]
@parametrize_dtype
def test_add_edges(idtype):
# homogeneous graph
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
u = 0
v = 1
g.add_edges(u, v)
assert g.device == F.ctx()
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 3
u = [0]
v = [1]
g.add_edges(u, v)
assert g.device == F.ctx()
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 4
u = F.tensor(u, dtype=idtype)
v = F.tensor(v, dtype=idtype)
g.add_edges(u, v)
assert g.device == F.ctx()
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 5
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))
# node id larger than current max node id
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
u = F.tensor([0, 1], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g.add_edges(u, v)
assert g.number_of_nodes() == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
# has data
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
u = F.tensor([0, 1], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g.add_edges(u, v, e_feat)
assert g.number_of_nodes() == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))
# zero data graph
g = dgl.graph([], num_nodes=0, idtype=idtype, device=F.ctx())
u = F.tensor([0, 1], dtype=idtype)
v = F.tensor([2, 2], dtype=idtype)
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g.add_edges(u, v, e_feat)
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 2
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([2, 2], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([2, 2], dtype=idtype))
assert F.array_equal(g.edata['hh'], F.tensor([2, 2], dtype=idtype))
# bipartite graph
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
u = 0
v = 1
g.add_edges(u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 3
u = [0]
v = [1]
g.add_edges(u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 4
u = F.tensor(u, dtype=idtype)
v = F.tensor(v, dtype=idtype)
g.add_edges(u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 5
u, v = g.edges(form='uv')
assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))
# node id larger than current max node id
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g.add_edges(u, v)
assert g.device == F.ctx()
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
# has data
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
g.ndata['h'] = {'user' : F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx()),
'game' : F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())}
g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
e_feat = {'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
'hh' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
g.add_edges(u, v, e_feat)
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_edges() == 4
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 0], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 0], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))
# heterogeneous graph
g = create_test_heterograph4(idtype)
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g.add_edges(u, v, etype='plays')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_nodes('developer') == 2
assert g.number_of_edges('plays') == 6
assert g.number_of_edges('develops') == 2
u, v = g.edges(form='uv', order='eid', etype='plays')
assert F.array_equal(u, F.tensor([0, 1, 1, 2, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 0, 1, 1, 2, 3], dtype=idtype))
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 0, 0], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 1, 1, 1, 0, 0], dtype=idtype))
# add with feature
e_feat = {'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
u = F.tensor([0, 2], dtype=idtype)
v = F.tensor([2, 3], dtype=idtype)
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 1, 1], dtype=idtype), ctx=F.ctx())
g.add_edges(u, v, data=e_feat, etype='develops')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 4
assert g.number_of_nodes('developer') == 3
assert g.number_of_edges('plays') == 6
assert g.number_of_edges('develops') == 4
u, v = g.edges(form='uv', order='eid', etype='develops')
assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 2, 3], dtype=idtype))
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3, 0], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 1, 1], dtype=idtype))
assert F.array_equal(g.edges['develops'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype))
@parametrize_dtype
def test_add_nodes(idtype):
# homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1,1,1], dtype=idtype), ctx=F.ctx())
g.add_nodes(1)
assert g.number_of_nodes() == 4
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
# zero node graph
g = dgl.graph([], num_nodes=3, idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1,1,1], dtype=idtype), ctx=F.ctx())
g.add_nodes(1, data={'h' : F.copy_to(F.tensor([2], dtype=idtype), ctx=F.ctx())})
assert g.number_of_nodes() == 4
assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 2], dtype=idtype))
# bipartite graph
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
g.add_nodes(2, data={'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, ntype='user')
assert g.number_of_nodes('user') == 4
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([0, 0, 2, 2], dtype=idtype))
g.add_nodes(2, ntype='game')
assert g.number_of_nodes('game') == 5
# heterogeneous graph
g = create_test_heterograph4(idtype)
g.add_nodes(1, ntype='user')
g.add_nodes(2, data={'h' : F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, ntype='game')
g.add_nodes(0, ntype='developer')
assert g.number_of_nodes('user') == 4
assert g.number_of_nodes('game') == 4
assert g.number_of_nodes('developer') == 2
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2, 2], dtype=idtype))
@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet has error with (0,) shape tensor.")
@parametrize_dtype
def test_remove_edges(idtype):
# homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
e = 0
g.remove_edges(e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
e = [0]
g.remove_edges(e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
e = F.tensor([0], dtype=idtype)
g.remove_edges(e)
assert g.number_of_edges() == 0
# has node data
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.ndata['h'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g.remove_edges(1)
assert g.number_of_edges() == 1
assert F.array_equal(g.ndata['h'], F.tensor([1, 2, 3], dtype=idtype))
# has edge data
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
g.remove_edges(0)
assert g.number_of_edges() == 1
assert F.array_equal(g.edata['h'], F.tensor([2], dtype=idtype))
# invalid eid
assert_fail = False
try:
g.remove_edges(1)
except:
assert_fail = True
assert assert_fail
# bipartite graph
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
e = 0
g.remove_edges(e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
e = [0]
g.remove_edges(e)
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
e = F.tensor([0], dtype=idtype)
g.remove_edges(e)
assert g.number_of_edges() == 0
# has data
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
g.ndata['h'] = {'user' : F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx()),
'game' : F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())}
g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
g.remove_edges(1)
assert g.number_of_edges() == 1
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2, 2], dtype=idtype))
assert F.array_equal(g.edata['h'], F.tensor([1], dtype=idtype))
# heterogeneous graph
g = create_test_heterograph4(idtype)
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
g.remove_edges(1, etype='plays')
assert g.number_of_edges('plays') == 3
u, v = g.edges(form='uv', order='eid', etype='plays')
assert F.array_equal(u, F.tensor([0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 1], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 3, 4], dtype=idtype))
# remove all edges of 'develops'
g.remove_edges([0, 1], etype='develops')
assert g.number_of_edges('develops') == 0
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2, 2], dtype=idtype))
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3], dtype=idtype))
@parametrize_dtype
def test_remove_nodes(idtype):
# homogeneous Graphs
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
n = 0
g.remove_nodes(n)
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
n = [1]
g.remove_nodes(n)
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 0
g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
n = F.tensor([2], dtype=idtype)
g.remove_nodes(n)
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
# invalid nid
assert_fail = False
try:
g.remove_nodes(3)
except:
assert_fail = True
assert assert_fail
# has node and edge data
g = dgl.graph(([0, 0, 2], [0, 1, 2]), idtype=idtype, device=F.ctx())
g.ndata['hv'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g.remove_nodes(F.tensor([0], dtype=idtype))
assert g.number_of_nodes() == 2
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
assert F.array_equal(g.ndata['hv'], F.tensor([2, 3], dtype=idtype))
assert F.array_equal(g.edata['he'], F.tensor([3], dtype=idtype))
# node id larger than current max node id
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
n = 0
g.remove_nodes(n, ntype='user')
assert g.number_of_nodes('user') == 1
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([2], dtype=idtype))
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
n = [1]
g.remove_nodes(n, ntype='user')
assert g.number_of_nodes('user') == 1
assert g.number_of_nodes('game') == 3
assert g.number_of_edges() == 1
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0], dtype=idtype))
assert F.array_equal(v, F.tensor([1], dtype=idtype))
g = dgl.bipartite(([0, 1], [1, 2]), 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
n = F.tensor([0], dtype=idtype)
g.remove_nodes(n, ntype='game')
assert g.number_of_nodes('user') == 2
assert g.number_of_nodes('game') == 2
assert g.number_of_edges() == 2
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
assert F.array_equal(v, F.tensor([0 ,1], dtype=idtype))
# heterogeneous graph
g = create_test_heterograph4(idtype)
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
g.remove_nodes(0, ntype='game')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 1
assert g.number_of_nodes('developer') == 2
assert g.number_of_edges('plays') == 2
assert g.number_of_edges('develops') == 1
assert F.array_equal(g.nodes['user'].data['h'], F.tensor([1, 1, 1], dtype=idtype))
assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2], dtype=idtype))
assert F.array_equal(g.nodes['developer'].data['h'], F.tensor([3, 3], dtype=idtype))
u, v = g.edges(form='uv', order='eid', etype='plays')
assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 0], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([3, 4], dtype=idtype))
u, v = g.edges(form='uv', order='eid', etype='develops')
assert F.array_equal(u, F.tensor([1], dtype=idtype))
assert F.array_equal(v, F.tensor([0], dtype=idtype))
if __name__ == '__main__':
# test_create()
......@@ -2082,7 +2495,7 @@ if __name__ == '__main__':
# test_inc()
# test_view("int32")
# test_view1("int32")
# test_flatten()
# test_flatten(F.int32)
# test_convert_bound()
# test_convert()
# test_to_device("int32")
......@@ -2101,6 +2514,10 @@ if __name__ == '__main__':
# test_bipartite()
# test_dtype_cast()
# test_reverse("int32")
test_reverse("int32")
test_format()
# test_format()
test_add_edges(F.int32)
test_add_nodes(F.int32)
test_remove_edges(F.int32)
test_remove_nodes(F.int32)
test_clone(F.int32)
pass
tests/compute/test_inplace_update.py → tests/compute/test_inplace_update.py.bak
View file @ 44089c8b
......@@ -21,57 +21,6 @@ def generate_graph():
return g
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
def test_inplace_recv():
u = F.tensor([0, 0, 0, 3, 4, 9])
v = F.tensor([1, 2, 3, 9, 9, 0])
def message_func(edges):
return {'m' : edges.src['f'] + edges.dst['f']}
def reduce_func(nodes):
return {'f' : F.sum(nodes.mailbox['m'], 1)}
def apply_func(nodes):
return {'f' : 2 * nodes.data['f']}
def _test(apply_func):
g = generate_graph()
f = g.ndata['f']
# one out place run to get result
g.send((u, v), message_func)
g.recv([0,1,2,3,9], reduce_func, apply_func)
result = g.get_n_repr()['f']
# inplace deg bucket run
v1 = F.clone(f)
g.ndata['f'] = v1
g.send((u, v), message_func)
g.recv([0,1,2,3,9], reduce_func, apply_func, inplace=True)
r1 = g.get_n_repr()['f']
# check result
assert F.allclose(r1, result)
# check inplace
assert F.allclose(v1, r1)
# inplace e2v
v1 = F.clone(f)
g.ndata['f'] = v1
g.send((u, v), message_func)
g.recv([0,1,2,3,9], fn.sum(msg='m', out='f'), apply_func, inplace=True)
r1 = g.ndata['f']
# check result
assert F.allclose(r1, result)
# check inplace
assert F.allclose(v1, r1)
# test send_and_recv with apply_func
_test(apply_func)
# test send_and_recv without apply_func
_test(None)
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
def test_inplace_snr():
u = F.tensor([0, 0, 0, 3, 4, 9])
......@@ -275,7 +224,7 @@ def test_inplace_apply():
# test apply all nodes, should not be done in place
g.ndata['f'] = nf
g.apply_nodes(apply_node_func, inplace=True)
assert F.allclose(nf, g.ndata['f']) == False
assert (F.allclose(nf, g.ndata['f']) == False)
edges = [3, 5, 7, 10]
ef = g.edata['e']
......@@ -293,7 +242,6 @@ def test_inplace_apply():
assert F.allclose(ef, g.edata['e']) == False
if __name__ == '__main__':
test_inplace_recv()
test_inplace_snr()
test_inplace_push()
test_inplace_pull()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment