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
  1. 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.


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.


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)


__init__([input, output])