probly.layers.flax

flax layer implementation.

Classes

BatchEnsembleConv	BatchEnsemble convolutional layer based on [WTB20].
BatchEnsembleLinear	BatchEnsemble Linear layer based on [WTB20].
DropConnectLinear	Custom Linear layer with DropConnect applied to weights during training based on [ASSR20].