pdd#

Framework-neutral projection, training, and sampling primitives for PDD.

This module owns projection layout and fusion plus the data-dependent and data-free objectives and block sampler. Model calls and architecture-specific packing remain behind the adapter protocol and belong in modelopt.torch.fastgen.plugins.

Classes

PDDLayerSpec

Immutable description of an architecture's final PDD projection.

PDDModelAdapter

Architecture adapter used by the framework-neutral PDD pipeline.

PDDOutputProjection

A widened linear projection with one output head per PDD interval.

PDDPipeline

PDD losses and fused sampler over a single core-owned grid.

Functions

convert_to_pdd_output_projection

Explicitly replace layer_spec.projection_path with a PDD projection.
class PDDLayerSpec#

Bases: object

Immutable description of an architecture’s final PDD projection.

channel_major stores widened outputs as [head, base_output]. patch_major stores them as [patch, head, output_channel] and therefore requires the unpatched output_channels count.

__init__(projection_path, head_layout, output_channels=None)#
Parameters:

  • projection_path (str)

  • head_layout (Literal['channel_major', 'patch_major'])

  • output_channels (int | None)

Return type:

None

head_layout: Literal['channel_major', 'patch_major']#
output_channels: int | None = None#
projection_path: str#
class PDDModelAdapter#

Bases: Protocol

Architecture adapter used by the framework-neutral PDD pipeline.

__init__(*args, **kwargs)#
student_all_heads(model, state, time, *, condition=None, **model_kwargs)#

Return canonical [batch, head, *latent_shape] student velocities.

Parameters:

  • model (Module)

  • state (Tensor)

  • time (Tensor)

  • condition (Any)

  • model_kwargs (Any)

Return type:

Tensor

student_fused_block(model, state, time, *, start, end, grid, condition=None, **model_kwargs)#

Return one base-shaped velocity from the fused projection block.

Parameters:

  • model (Module)

  • state (Tensor)

  • time (Tensor)

  • start (int)

  • end (int)

  • grid (Tensor)

  • condition (Any)

  • model_kwargs (Any)

Return type:

Tensor

teacher_velocity(model, state, time, *, condition=None, negative_condition=None, **model_kwargs)#

Return the adapter-specific guided teacher velocity.

Parameters:

  • model (Module)

  • state (Tensor)

  • time (Tensor)

  • condition (Any)

  • negative_condition (Any)

  • model_kwargs (Any)

Return type:

Tensor

class PDDOutputProjection#

Bases: Linear

A widened linear projection with one output head per PDD interval.

forward returns the full widened output unless an explicit fusion tuple (start, end, grid) is supplied. Fused parameters are computed in float32 and applied without mutating the module or replacing its registered weights.

__init__(in_features, base_out_features, grid_size, layer_spec, *, bias=True, device=None, dtype=None)#

Initialize an unpopulated widened projection with validated layout metadata.

Parameters:

  • in_features (int)

  • base_out_features (int)

  • grid_size (int)

  • layer_spec (PDDLayerSpec)

  • bias (bool)

  • device (device | str | None)

  • dtype (dtype | None)

Return type:

None

forward(input, *, fusion=None)#

Apply the widened projection or an explicitly selected fused block.

Parameters:

  • input (Tensor)

  • fusion (tuple[int, int, Tensor] | None)

Return type:

Tensor

classmethod from_linear(base_linear, grid_size, layer_spec)#

Convert a loaded base linear without modifying it.

Repeating a compatible conversion is idempotent. A PDD projection with conflicting grid, layout, path, or channel metadata is rejected.

Parameters:

  • base_linear (Linear)

  • grid_size (int)

  • layer_spec (PDDLayerSpec)

Return type:

PDDOutputProjection

property patch_factor: int#

Number of output patches represented by the base projection.

class PDDPipeline#

Bases: object

PDD losses and fused sampler over a single core-owned grid.

__init__(student, teacher, config, adapter)#

Store the models/config/adapter and freeze the optional training teacher.

Parameters:
Return type:

None

compute_data_free_loss(state, *, n, condition=None, negative_condition=None, model_kwargs=None, k=None, generator=None, collect_metrics=True)#

Compute Algorithm 3 and return the detached state carried to n + L_min.

Parameters:

  • state (Tensor)

  • n (Tensor)

  • condition (Any)

  • negative_condition (Any)

  • model_kwargs (Mapping[str, Any] | None)

  • k (Tensor | None)

  • generator (Generator | None)

  • collect_metrics (bool)

Return type:

tuple[Tensor, dict[str, Tensor], Tensor, Tensor]

compute_loss(data, *, noise=None, condition=None, negative_condition=None, model_kwargs=None, n=None, k=None, generator=None, collect_metrics=True)#

Compute the exact data-dependent PDD objective for one batch.

Parameters:

  • data (Tensor)

  • noise (Tensor | None)

  • condition (Any)

  • negative_condition (Any)

  • model_kwargs (Mapping[str, Any] | None)

  • n (Tensor | None)

  • k (Tensor | None)

  • generator (Generator | None)

  • collect_metrics (bool)

Return type:

tuple[Tensor, dict[str, Tensor]]

sample(noise, *, condition=None, blocks=None, model_kwargs=None)#

Sample from raw RF noise with one fused call per contiguous block.

Parameters:

  • noise (Tensor)

  • condition (Any)

  • blocks (Sequence[int] | None)

  • model_kwargs (Mapping[str, Any] | None)

Return type:

Tensor

time_grid(device=None)#

Construct this pipeline’s sole shifted rectified-flow grid.

Parameters:

device (device | str | None)

Return type:

Tensor

convert_to_pdd_output_projection(model, layer_spec, grid_size)#

Explicitly replace layer_spec.projection_path with a PDD projection.

Parameters:
Return type:

PDDOutputProjection