Carlo Lucibello

- [ ] ForwardDiff - [ ] NNPACK_jll

Extend with the `dims` keyword similarly to https://pytorch.org/docs/stable/generated/torch.gather.html#torch-gather TODO - [ ] docs - [ ] more tests - [ ] gradient - [ ] non-mutating version

See for example hugging face datasets repo structure https://huggingface.co/docs/datasets/repository_structure and torchvision [DatasetFolder](https://pytorch.org/vision/stable/generated/torchvision.datasets.DatasetFolder.html#torchvision.datasets.DatasetFolder.find_classes)

fix #191 Still lot to do

fix #184

Enzyme.jl compatibility

8

Here we will keep track of compatibility with Enzyme for taking gradients. First think is to collect a few examples to run.

Support Metal

6

Needs gather/scatter kernels in NNlib

Layers support for HeteroGraphConv

6

HeteroGraphConv are build out of standard graph conv layers which are individually applied to the different relations. The list of layers supporting integration with HeteroGraphConv should be extended. - [...

convenience feature setter

1

Right now we can conveniently add new features with ```julia julia> g = rand_graph(3, 2) GNNGraph: num_nodes: 3 num_edges: 2 julia> g.ndata.x = rand(1, 3) 1×3 Matrix{Float64}: 0.31553 0.824744 0.86907...

Lux support?

2

This package is built on top of Flux.jl, but we could consider also supporting Lux.jl since it is being used in some part of the ecosystem. It is not clear...