neural_tangents.stax.Diagonal
- class neural_tangents.stax.Diagonal(input=Bool.YES, output=Bool.NO)[source]
Helps decide whether to allow the kernel to contain diagonal entries only.
- The intended behavior is to be diagonal-only iff
output off-diagonal entries are all zeros, and
b) diagonal-only
Kernel
is sufficient for all steps of computation.
Note that currently this parameter is shared between all parallel branches, even if this is excessive, and it is defined once for the whole network and does not change from layer to layer, even if it could be possible.
- Must be endowed with
1) A commutative, associative, idempotent
AND
(&
) operation, corresponding to combining requirements of two layers in parallel.2) An associative composition
>>
operation, corresponding to the requirement of a composition of two layers.
- input
specifies whether inputs to given layer can contain only diagonal entries.
Bool.YES
means “yes”;Bool.MAYBE
means iff off-diagonal entries are zero.Bool.NO
means “no”. When traversing the network tree from inputs to outputs (as well as parallel branches from left/right to right/left) can only decrease.
- output
specifies whether any outputs (starting from this layer to the output of the network) can contain only diagonal entries.
Bool.YES
means yes;Bool.MAYBE
means “yes” after current layer, but may become “no” further in the network.Bool.NO
means “no”.
- __init__(input=Bool.YES, output=Bool.NO)
Methods
__init__
([input, output])Attributes