neural_tangents.stax.AggregateImplementation

class neural_tangents.stax.AggregateImplementation(value)[source]

Implementation of the Aggregate layer.

See Aggregate docstring for details.

DENSE

Is recommended for dense graphs, where the number of edges E is proportional to the number of vertices V to the power of 1.5 or more.

SPARSE

Is recommended for sparse graphs, where E ~ O(V) or less.

__init__()

Attributes

DENSE

SPARSE