base

Base module for noise layers.

Classes:

Name	Description
BaseNoiseLayer	Applies a stochastic transformation to a Tensor using the given Estimators.

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, 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

Methods:

Name	Description
__call__	Transform the input data.
__getstate__	Prepare a copy of the noise layer's state for serialization.
__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, and the generator
dict[str, Any]	states if they exist.

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

__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 and _generators are both 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.

forward

forward(
    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 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.

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.

Parameters:

Name	Type	Description	Default
seed	int | None	The seed to set.	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.