base

Base module for noise layers.

Classes:

Name	Description
`BaseNoiseLayer`	Applies a stochastic transformation to a `Tensor` using the given `Estimator`s.

Attributes:

Name	Type	Description
`POST_PROCESSOR_STATE_KEY`	`Final[str]`	The key under which a serialized noise layer records its attached post-processor's type and configuration.

POST_PROCESSOR_STATE_KEY `module-attribute` ¶

POST_PROCESSOR_STATE_KEY: Final[str] = 'post_processor'

The key under which a serialized noise layer records its attached post-processor's type and configuration.

BaseNoiseLayer ¶

Bases: Module, ABC, Generic[EstimatorModuleT, ParameterizationT, OptionalMaskerT_contra]

Applies a stochastic transformation to a Tensor using the given Estimators.

Parameters:

Name	Type	Description	Default
`seed` ¶	`int \| None`	Seed for the random number generator used to generate the stochastic transformation. If `None`, the global RNG state is used.	required
`mean_estimator` ¶	`Estimator[EstimatorModuleT, BoundedMeanParameterization \| None, None]`	The estimator to use to estimate the mean of the stochastic transformation.	required
`std_estimator` ¶	`Estimator[EstimatorModuleT, ParameterizationT, OptionalMaskerT_contra]`	The estimator to use to estimate the standard deviation and optional input mask of the stochastic transformation.	required
`noise_layer_dtype` ¶	`dtype \| None`	The `torch.dtype` to use for the noise layer. If `None`, the default dtype is used.	`None`
`remove_skip_connection` ¶	`bool`	Whether to remove the skip connection in the noise layer. This is highly experimental and should be used with caution.	`False`
`post_processor` ¶	`PostProcessor \| None`	An optional `PostProcessor` to apply to the composed output of the stochastic transformation. It becomes part of this noise layer's module tree, so its weights travel in this noise layer's state dictionary and its configuration is nested inside this noise layer's serialized configuration. Its parameters must be in `noise_layer_dtype`, like every other parameter of the noise layer.	`None`

Methods:

Name	Description
`__call__`	Transform the input data.
`__getstate__`	Prepare a copy of the noise layer's state for serialization.
`__init__`
`__setstate__`	Restore the state of the noise layer from a serialized copy.
`forward`	Transform the input data.
`get_applied_transform_components_factory`	Create a function that returns the elements of the transform components (`'mean'` and `'std'`) applied during the most recent
`get_transformed_output_factory`	Create a function that returns the transformed output from the most recent forward pass.
`initial_seed`	Return the initial seed of the CPU device's random number generator.
`manual_seed`	Seed each of the random number generators.
`reset_parameters`	Reinitialize parameters and buffers.
`seed`	Seed each of the random number generators using a non-deterministic random number.

call ¶

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> torch.Tensor

Transform the input data.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask that selects the elements of `input` to transform. Where the mask is `False`, the original `input` value is returned. Also used to select the elements of the sampled standard deviations to use to mask the `input`. If `None`, the entire `input` is transformed.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to the estimator modules.	required

getstate ¶

__getstate__() -> dict[str, Any]

Prepare a copy of the noise layer's state for serialization.

The returned dictionary can be JSON-serialized if the state_dict key is removed (it contains tensors). Generally you should pop the state_dict, save it as a safetensors file, then JSON-serialize the rest of the dictionary.

Returns:

Type	Description
`dict[str, Any]`	A dictionary containing the configuration of the noise layer, including its type string, the state dict, the post-processor
`dict[str, Any]`	configuration if one is attached, and the generator states if they exist.

Changed in version v3.15.0: Added serialization support for all noise layers.

Changed in version v3.78.0: An attached `post_processor`'s type and configuration are now included, under the `post_processor` key. Its weights are already covered by the `state_dict` key.

init ¶

__init__(
    seed: int | None,
    mean_estimator: Estimator[
        EstimatorModuleT,
        BoundedMeanParameterization | None,
        None,
    ],
    std_estimator: Estimator[
        EstimatorModuleT,
        ParameterizationT,
        OptionalMaskerT_contra,
    ],
    noise_layer_dtype: dtype | None = None,
    remove_skip_connection: bool = False,
    *,
    post_processor: PostProcessor | None = None
) -> None

Changed in version v3.32.0: Added `remove_skip_connection` argument to optionally remove the skip connection in the transformer cloak architecture. This is highly experimental and should be used with caution.

Changed in version v3.78.0: Added the optional `post_processor` argument, which applies a learned transformation to the composed output of the stochastic transformation.

setstate ¶

__setstate__(
    state: dict[str, Any],
    trust_remote_code: bool = False,
    third_party_model_path: (
        str | PathLike[str] | None
    ) = None,
) -> None

Restore the state of the noise layer from a serialized copy.

state_dict, post_processor, and _generators are all optional keys, and will be restored if they exist in the state.

Parameters:

Name	Type	Description	Default
`state` ¶	`dict[str, Any]`	The state to set.	required
`trust_remote_code` ¶	`bool`	Whether to trust remote code when loading from HuggingFace Hub.	`False`
`third_party_model_path` ¶	`str \| PathLike[str] \| None`	The path or huggingface reference to a third-party model to load. This is useful when loading SGTs whose internal structure depends on transformers which are not importable directly through transformers, but are present on the Hugging Face Hub.	`None`

Changed in version v3.15.0: Added serialization support for all noise layers.

Changed in version v3.78.0: A post-processor recorded under the `post_processor` key is now reconstructed and passed to the constructor.

forward ¶

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

Transform the input data.

An attached post_processor is applied last, to the composed output. Subclasses that further transform the value returned by this method (for example, by clamping it to a valid pixel range) therefore apply their own transform after the post-processor.

Note

Because the noise mask is applied to this method's return value by a forward hook (see __call__), a post-processor only ever affects the positions the noise layer transformed; elsewhere the original input value is retained regardless of what the post-processor computed.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask that selects the elements of `input` to transform. Where the mask is `0`, the original `input` value is returned. Also used to select the elements of the sampled standard deviations to use to mask the `input`. If `None`, the entire `input` is transformed.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to the estimator modules.	required

Returns:

Type	Description
`torch.Tensor`	The transformed input data.

Changed in version v3.64.0: The skip connection term is now taken from the overridable `_skip_connection_input` hook, which defaults to returning the input unchanged. Behavior is unchanged for every layer that does not override it.

Changed in version v3.78.0: An attached `post_processor` is now applied to the composed output. Behavior is unchanged for every noise layer without one, which is the default.

get_applied_transform_components_factory ¶

get_applied_transform_components_factory() -> (
    Callable[[], dict[str, torch.Tensor]]
)

Create a function that returns the elements of the transform components ('mean' and 'std') applied during the most recent forward pass.

Specifically, the applied elements are those selected by the noise mask (if supplied) and standard deviation mask (if std_estimator.masker is not None). If no masks are used, all elements are returned.

The applied transform components are returned flattened.

This function is intended to be used to log histograms of the transform components.

Returns:

Type	Description
`Callable[[], dict[str, torch.Tensor]]`	A function that returns the the elements of the transform components applied during the most recent forward pass.

Examples:

>>> from torch import nn
>>> from stainedglass_core import model as sg_model, noise_layer as sg_noise_layer
>>> base_model = nn.Linear(20, 2)
>>> noisy_model = sg_model.NoisyModel(
...     sg_noise_layer.CloakNoiseLayer1,
...     base_model,
...     target_parameter="input",
... )
>>> get_applied_transform_components = (
...     noisy_model.noise_layer.get_applied_transform_components_factory()
... )
>>> input = torch.ones(1, 20)
>>> noise_mask = torch.tensor(5 * [False] + 15 * [True])
>>> output = noisy_model(input, noise_mask=noise_mask)
>>> applied_transform_components = get_applied_transform_components()
>>> applied_transform_components
{'mean': tensor(...), 'std': tensor(...)}
>>> {
...     component_name: component.shape
...     for component_name, component in applied_transform_components.items()
... }
{'mean': torch.Size([15]), 'std': torch.Size([15])}

get_transformed_output_factory ¶

get_transformed_output_factory() -> (
    Callable[[], torch.Tensor]
)

Create a function that returns the transformed output from the most recent forward pass.

If super batching is active, only the transformed half of the super batch output is returned.

Returns:

Type	Description
`Callable[[], torch.Tensor]`	A function that returns the transformed output from the most recent forward pass.

Examples:

>>> from stainedglass_core import noise_layer as sg_noise_layer
>>> noise_layer = sg_noise_layer.CloakNoiseLayer1()
>>> get_transformed_output = noise_layer.get_transformed_output_factory()
>>> input = torch.ones(2, 3, 32, 32)
>>> output = noise_layer(input)
>>> transformed_output = get_transformed_output()
>>> assert output.equal(transformed_output)

initial_seed ¶

initial_seed() -> int

Return the initial seed of the CPU device's random number generator.

manual_seed ¶

manual_seed(
    seed: int | None, rank_dependent: bool = True
) -> None

Seed each of the random number generators.

Setting seed to None will destroy any existing generators, disabling reproducible behavior until the next call with an integer seed. Calling manual_seed again with an integer after a None call fully recreates generators for every currently available device type (CPU, CUDA, MPS) before applying the seed, making the operation safely reversible.

Parameters:

Name	Type	Description	Default
`seed` ¶	`int \| None`	The seed to set. Pass `None` to disable reproducible behavior.	required
`rank_dependent` ¶	`bool`	Whether to add the distributed rank to the seed to ensure that each process samples different noise.	`True`

reset_parameters ¶

reset_parameters() -> None

Reinitialize parameters and buffers.

This method is useful for initializing tensors created on the meta device.

seed ¶

seed() -> None

Seed each of the random number generators using a non-deterministic random number.

base

POST_PROCESSOR_STATE_KEY `module-attribute` ¶

BaseNoiseLayer ¶

`seed` ¶

`mean_estimator` ¶

`std_estimator` ¶

`noise_layer_dtype` ¶

`remove_skip_connection` ¶

`post_processor` ¶

call ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

getstate ¶

init ¶

setstate ¶

`state` ¶

`trust_remote_code` ¶

`third_party_model_path` ¶

forward ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

get_applied_transform_components_factory ¶

get_transformed_output_factory ¶

initial_seed ¶

manual_seed ¶

`seed` ¶

`rank_dependent` ¶

reset_parameters ¶

seed ¶

base

POST_PROCESSOR_STATE_KEY module-attribute ¶

BaseNoiseLayer ¶

seed ¶

mean_estimator ¶

std_estimator ¶

noise_layer_dtype ¶

remove_skip_connection ¶

post_processor ¶

__call__ ¶

input ¶

noise_mask ¶

**kwargs ¶

__getstate__ ¶

__init__ ¶

__setstate__ ¶

state ¶

trust_remote_code ¶

third_party_model_path ¶

forward ¶

input ¶

noise_mask ¶

**kwargs ¶

get_applied_transform_components_factory ¶

get_transformed_output_factory ¶

initial_seed ¶

manual_seed ¶

seed ¶

rank_dependent ¶

reset_parameters ¶

seed ¶

POST_PROCESSOR_STATE_KEY `module-attribute` ¶

`seed` ¶

`mean_estimator` ¶

`std_estimator` ¶

`noise_layer_dtype` ¶

`remove_skip_connection` ¶

`post_processor` ¶

call ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

getstate ¶

init ¶

setstate ¶

`state` ¶

`trust_remote_code` ¶

`third_party_model_path` ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

`seed` ¶

`rank_dependent` ¶