Not possible to sample twice along same edge_set?

[cgparkinson]

Hi,

I have a heterogeneous graph where I would like to sample twice along the same edge set, i.e. starting from A we can reach B, C, D, E, F, G:

  B - C
 / \
A - D - E
 \
  F - G

This should be possible, according to the documentation:

# Store builder at level of "author written papers":
builder = tfgnn.SamplingSpecBuilder(schema).seed('author').sample(10, 'writes')
path1 = builder.sample(5, 'cited_by')
path2 = builder.sample(3, 'written_by').sample('writes')

proto = (tfgnn.SamplingSpecBuilder.join([path1, path2]).sample(10, 'cited_by')
         .build())

# The above `Join()` can be made less verbose with:
proto = path1.Join([path2]).sample(10, 'cited_by').build()

In the above example, cited_by is clearly used twice.

However when I try this I get the error message:

ValueError: The name ... is used 2 times in the model. All layer names should be unique.

To reproduce:

I can easily reproduce this in the ogbn notebook, by replacing:

  papers_cited_from_seed = seed.sample(sizes["cites"], "cites")

with

  papers_cited_from_seed = seed.sample(sizes["cites"], "cites").sample(sizes["cites"], "cites")

or equivalently

  papers_cited_from_seed = seed.sample([sizes["cites"], sizes["cites"]], "cites")

to obtain the error:

ValueError: The name "cites" is used 2 times in the model. All layer names should be unique.

Is this a bug or is this intentional? It's possible I'm thinking along the wrong lines.

Many thanks.