ls_mlkit.diffuser.base_diffuser module

Base Diffuser Config and Base Diffuser.

class ls_mlkit.diffuser.base_diffuser.BaseDiffuser(config: BaseDiffuserConfig, time_scheduler: TimeScheduler)[source]

Bases: Module, ABC

Base Diffuser Class

complete_micro_shape(x: Tensor) → Tensor[source]

Complete the micro shape of \(x\), assuming the macro shape is already known

Parameters:

x (Tensor) – \(x\)

Returns:

\(x\) with the micro shape completed

Return type:

Tensor

abstractmethod compute_loss(batch: dict[str, Any], *args: list[Any], **kwargs: dict[Any, Any]) → dict[source]

Compute loss

Parameters:

batch (dict[str, Any]) – the batch of data

Returns:

a dictionary that must contain the key “loss”

Return type:

dict

forward(batch: dict[str, Any], *args: list[Any], **kwargs: dict[Any, Any])[source]

Forward function, input batch of data and return the dictionary containing the loss

Parameters:

batch (dict[str, Any]) – the batch of data

Returns:

a dictionary that must contain the key “loss”

Return type:

dict

abstractmethod forward_process(x_0: Tensor, discrete_t: Tensor, mask: Tensor, *args: list[Any], **kwargs: dict[Any, Any]) → dict[source]

Forward process, from \(x_0\) to \(x_t\)

Parameters:

  • x_0 (Tensor) – \(x_0\)

  • discrete_t (Tensor) – the discrete time steps \(t\)

  • mask (Tensor) – the mask of the sample

Returns:

a dictionary that must contain the key “x_t”

Return type:

dict

get_macro_and_micro_shape(x: Tensor) → Tuple[Tuple[int, ...], Tuple[int, ...]][source]

Get the macro and micro shape of \(x\)

Parameters:

x (Tensor) – \(x\)

Returns:

the macro and micro shape of \(x\)

Return type:

Tuple[Tuple[int, ...], Tuple[int, ...]]

get_macro_shape(x: Tensor) → Tuple[int, ...][source]

Get the macro shape of \(x\)

Parameters:

x (Tensor) – \(x\)

Returns:

the shape of the macro part of \(x\)

Return type:

Tuple[int, ...]

abstractmethod prior_sampling(shape: Tuple[int, ...]) → Tensor[source]

Sample initial noise used for reverse process

Parameters:

shape (Tuple[int, ...]) – the shape of the sample

Returns:

the initial noise

Return type:

Tensor

abstractmethod sample_xtm1_conditional_on_xt(x_t, discrete_t: Tensor, padding_mask: Tensor, *args: list[Any], **kwargs: dict[Any, Any]) → Tensor[source]

Sample \(x_{t-1}\) conditional on \(x_t\)

Parameters:

  • x_t – \(x_t\)

  • discrete_t (Tensor) – the discrete time steps \(t\)

  • padding_mask (Tensor) – the padding mask

Returns:

\(x_{t-1}\)

Return type:

Tensor

class ls_mlkit.diffuser.base_diffuser.BaseDiffuserConfig(n_discretization_steps: int, ndim_micro_shape: int, *args: list[Any], **kwargs: dict[Any, Any])[source]

Bases: BaseConfigClass

Diffuser configure base class