post_processors

Module for noise layer post-processors.

A post-processor is an optional, learnable module that runs on the output of a noise layer's stochastic composition, i.e. on mean(x) + std(x) * N(0, 1) rather than on any of its parts. A noise layer that supports post-processors accepts one through the post_processor argument of its constructor and owns it as a submodule: its weights travel in the noise layer's state dictionary, and its type and configuration are recorded in the noise layer's serialized state, from which deserialization reconstructs it and passes it back to the noise layer's constructor.

Because the noise mask is applied to the noise layer's output after forward returns (see BaseNoiseLayer.__call__), a post-processor only ever affects the positions the noise layer actually transformed. Positions where the noise mask is False retain their original, untransformed input value regardless of what the post-processor computes for them.

This module defines the interface, including how a post-processor is serialized and deserialized. Concrete post-processors are defined alongside it.

Classes:

Name	Description
PostProcessor	Defines an interface for transforming the output of a noise layer's stochastic composition.

Attributes:

Name	Type	Description
POST_PROCESSOR_TYPE_KEY	Final[str]	The key under which a serialized post-processor records its fully qualified class name.

POST_PROCESSOR_TYPE_KEY module-attribute

POST_PROCESSOR_TYPE_KEY: Final[str] = 'post_processor_type'

The key under which a serialized post-processor records its fully qualified class name.

PostProcessor

Bases: Module, ABC

Defines an interface for transforming the output of a noise layer's stochastic composition.

Post-processors are shape-preserving: the tensor they return must be broadcastable to the tensor they were given, since the noise layer's noise mask is applied to their output against the noise layer's own input.

Methods:

Name	Description
deserialize	Reconstruct a post-processor from the dictionary produced by serialize.
forward	Transform the output of a noise layer's stochastic composition.
get_config	Return the configuration containing sufficient information to re-initialize this post-processor.
reset_parameters	Reinitialize parameters and buffers.
serialize	Serialize this post-processor's type and configuration to a JSON-serializable dictionary.

deserialize classmethod

deserialize(state: Mapping[str, Any]) -> Self

Reconstruct a post-processor from the dictionary produced by serialize.

Calling this on PostProcessor accepts any post-processor, while calling it on a subclass accepts only that subclass or its descendants.

The reconstructed post-processor's parameters hold whatever the class initializes them to; the containing noise layer is responsible for loading the trained weights from its own state dictionary.

Parameters:

Name	Type	Description	Default
state	Mapping[str, Any]	The serialized post-processor type and configuration.	required

Returns:

Type	Description
Self	The reconstructed post-processor.

Raises:

Type	Description
TypeError	If the serialized type is not a subclass of the class this is called on.

forward abstractmethod

forward(
    input: Tensor, noise_mask: Tensor | None = None
) -> torch.Tensor

Transform the output of a noise layer's stochastic composition.

Parameters:

Name	Type	Description	Default
input	Tensor	The noise layer's composed output, of shape (..., dim).	required
noise_mask	Tensor | None	An optional mask marking the elements of input that the noise layer transformed. Post-processors that operate pointwise can ignore this, since the noise layer discards their output wherever the mask is False anyway. It is supplied for post-processors that mix information across the masked dimension and therefore need to know which positions carry transformed values.	None

Returns:

Type	Description
torch.Tensor	The transformed input, of the same shape as input.

get_config abstractmethod

get_config() -> Mapping[str, Any]

Return the configuration containing sufficient information to re-initialize this post-processor.

The state dictionary must not be included, since it is serialized by the containing noise layer.

This must return a JSON-serializable mapping, since it is embedded in the noise layer's serialized configuration. The post_processor_type key is reserved for the serialized post-processor's fully qualified class name and must not be included.

reset_parameters

reset_parameters() -> None

Reinitialize parameters and buffers.

This method is useful for initializing tensors created on the meta device, which is how a post-processor is materialized when a noise layer is deserialized. It may be called while the parameters are still on the meta device and must not read their values. The default implementation does nothing, which is correct only for parameterless post-processors.

serialize

serialize() -> dict[str, Any]

Serialize this post-processor's type and configuration to a JSON-serializable dictionary.

Returns:

Type	Description
dict[str, Any]	This post-processor's fully qualified class name alongside its configuration.

Raises:

Type	Description
ValueError	If this post-processor's configuration uses the reserved post_processor_type key.