azula.nn.utils

Utilities for modules and networks.

Classes

FlattenWrapper

Creates a flatten/unflatten wrapper around a backbone.

Descriptions

class azula.nn.utils.FlattenWrapper(wrappee, shape=None)

Creates a flatten/unflatten wrapper around a backbone.

The purpose of this module is to create a flatten/unflatten frontier between azula components that opperate on one-dimensional vectors and a backbone that opperates on multi-dimensional tensors, like \(C \times H \times W\) images.

Parameters:

  • wrappee (Module) – The wrapped backbone.

  • shape (Sequence[int] | None) – The tensor shape.

forward(x_t, t, shape=None, **kwargs)
Parameters:

  • x_t (Tensor) – A noisy vector \(x_t\), with shape \((*, D)\).

  • t (Tensor) – The time \(t\), with shape \((*)\).

  • shape (Sequence[int] | None) – The unflattened tensor shape. If None, use self.shape instead.

  • kwargs – Optional keyword arguments.

Returns:

The output vector(s), with shape \((*, D)\).

Return type:

Tensor | List[Tensor]