Module contents

Module contents

Kaiwu-Pytorch-Plugin public API.

class kaiwu.torch_plugin.BoltzmannMachine(num_nodes: int, quadratic_coef: FloatTensor | None = None, linear_bias: FloatTensor | None = None, device=None)

Bases: AbstractBoltzmannMachine

Boltzmann Machine.

Args:

num_nodes (int): Total number of nodes in the model.

quadratic_coef (torch.FloatTensor, optional): quadratic coefficent,

shape is [num_nodes, num_nodes]

linear_bias (torch.FloatTensor, optional): linear bias, shape is [num_nodes]

device (torch.device, optional): Device for tensor construction. If None, uses CPU.

clip_parameters(h_range, j_range) → None

Clip linear and quadratic bias weights in-place.

Args:

h_range (tuple[float, float]): Range for quadratic weights. for example, [-1, 1] j_range (tuple[float, float]): Range for linear weights. for example, [-1, 1]

condition_sample(sampler, s_visible, dtype=torch.float32) → Tensor

Sample from the Boltzmann Machine given some nodes.

Args:

sampler (kaiwu.core.Optimizer): Optimizer used for sampling from the model. s_visible: State of the visible layer.

Returns:

torch.Tensor: Spins sampled from the model

(shape determined by sampler and sample_params).

forward(s_all: Tensor) → Tensor

Compute the Hamiltonian.

Args:

s_all (torch.tensor): Tensor of shape (B, N), where B is batch size,

N is the number of variables in the model.

Returns:

torch.tensor: Hamiltonian of shape (B,).

gibbs_sample(num_steps: int = 100, s_visible: Tensor | None = None, num_sample=None) → Tensor

Sample from the Boltzmann Machine.

Args:

num_steps (int): Number of Gibbs sampling steps.

s_visible (torch.Tensor, optional): State of the visible layer,

shape (B, num_visible). If None, randomly initialize visible layer.

num_sample (int, optional): Number of samples.

If None, uses batch size of s_visible.

hidden_bias(num_hidden: int) → Tensor

Get the hidden bias.

Args:

num_hidden (int): Number of hidden nodes.

symmetrized_quadratic_coef()

Quadratic coefficient

visible_bias(num_visible) → Tensor

Get the visible bias.

Args:

num_visible (int): Number of visible nodes.

class kaiwu.torch_plugin.EnergyModel(bm_num_visible: int | None = None, bm_num_hidden: int | None = None, sampler: Any | None = None)

Bases: Module

Energy-side model interface used by QDiffusion candidate scoring.

discretize_visible_state(visible_logits: Tensor) → Tensor

Converts visible logits into normalized BM visible conditions.

forward(noisy_tokens: Tensor, candidate_tokens: Tensor, attention_mask: Tensor) → Tensor | None

Runs conditioned energy scoring through the PyTorch module API.

get_last_stats() → dict[str, Tensor]

Returns lightweight sampler diagnostics from the last score call.

sample_hidden_state(visible_state: Tensor) → tuple[Tensor, Tensor]

Samples BM hidden states for each visible assignment.

score_conditioned(noisy_tokens: Tensor, candidate_tokens: Tensor, attention_mask: Tensor) → Tensor | None

Scores candidates conditioned on noisy tokens.

score_visible_logits(visible_logits: Tensor) → Tensor

Scores visible logits under the conditioned BM energy model.

class kaiwu.torch_plugin.QDiffusion(proposal_model: Module, energy_model: EnergyModel, token_spec: SequenceTokenSpec, config: QDiffusionConfig | None = None, dtype: dtype = torch.float32, device: device | str | None = None, freeze_proposal: bool = True)

Bases: Module

Energy-guided discrete diffusion wrapper over generic sequence backbones. Initializes a QDiffusion model.

The class combines two backbone roles:

• a proposal model that predicts token logits for the current noisy state

• an energy model that reranks candidate reconstructions

It exposes both training-oriented APIs such as objective() and decoding-oriented APIs such as initialize_state(), step(), and generate().

Args:

proposal_model: Backbone used to predict proposal logits.

energy_model: Energy-side model used to encode and score candidates.

token_spec: Special-token metadata required by the generator.

config: Optional generation/training configuration.

dtype: Floating point dtype tracked by the wrapper.

device: Optional target device. When omitted, infer from parameters.

freeze_proposal: Whether to freeze proposal model parameters.

energy(noisy_tokens: Tensor, candidate_tokens: Tensor, attention_mask: Tensor | None = None) → Tensor

Scores candidate reconstructions conditioned on the noisy state.

Args:

noisy_tokens: Noisy token tensor used as conditioning input.

candidate_tokens: Candidate clean token tensor to score.

attention_mask: Optional attention mask for the energy model.

Returns:

torch.Tensor: A tensor of scalar energies with shape [batch, 1].

forward(noisy_tokens: Tensor, **kwargs: Any) → Tensor

Runs the proposal model on the current noisy state.

Args:

noisy_tokens: Current noisy token tensor. **kwargs: Additional keyword arguments forwarded to the proposal

model.

Returns:

torch.Tensor: Proposal logits over the token vocabulary.

Raises:

TypeError: If the proposal model does not implement forward.

generate(input_tokens: Tensor, *, max_steps: int = 500, partial_masks: Tensor | None = None, temperature: float = 1.0, return_state: bool = False) → Tensor | dict[str, Any]

Runs a complete iterative decoding loop inside the core class.

Args:

input_tokens: Initial token tensor. max_steps: Number of decode iterations to run. partial_masks: Optional boolean mask of fixed positions. temperature: Sampling temperature stored in the decode state. return_state: Whether to return the full final state dictionary.

Returns:

torch.Tensor | dict[str, Any]: Either the final token tensor or the full decode state.

get_non_special_symbol_mask(output_tokens: Tensor, partial_masks: Tensor | None = None) → Tensor

Returns a boolean mask of editable non-special-token positions.

Args:

output_tokens: Token tensor to inspect. partial_masks: Optional boolean mask of positions that should remain

fixed.

Returns:

torch.Tensor: A boolean mask where True marks editable non-special positions.

initialize_state(input_tokens: Tensor, partial_masks: Tensor | None = None, max_steps: int = 500, temperature: float = 1.0) → dict[str, Any]

Creates the initial decoding state for an external generation loop.

Args:

input_tokens: Initial token tensor. partial_masks: Optional boolean mask of fixed positions. max_steps: Planned number of decode iterations. temperature: Sampling temperature stored in the state payload.

Returns:

dict[str, Any]: A mutable state dictionary suitable for repeated step() calls.

objective(batch: dict[str, Tensor], weighting: str = 'constant') → dict[str, Tensor]

Builds the one-step training objective used by an external loop.

Args:

batch: Batch dictionary containing at least batch["targets"]. weighting: Per-sample timestep weighting mode.

Returns:

dict[str, torch.Tensor]: A dictionary containing proposal logits, supervision masks, loss weights, and the EBM objective term.

proposal(noisy_tokens: Tensor, **kwargs: Any) → Tensor

Semantic alias around forward() for proposal-side calls.

Args:

noisy_tokens: Current noisy token tensor. **kwargs: Additional keyword arguments forwarded to the proposal

model.

Returns:

torch.Tensor: Proposal logits over the token vocabulary.

step(state: dict[str, Any], partial_masks: Tensor | None = None) → dict[str, Any]

Runs one denoising/reranking step and returns updated state.

Args:

state: Current decode state created by initialize_state(). partial_masks: Optional boolean mask of fixed positions.

Returns:

dict[str, Any]: The updated decode state after one iteration.

to(*args: Any, **kwargs: Any) → QDiffusion

Moves the module and refreshes cached device/dtype metadata.

Args:

*args: Positional arguments forwarded to nn.Module.to. **kwargs: Keyword arguments forwarded to nn.Module.to.

Returns:

QDiffusion: The moved module instance.

class kaiwu.torch_plugin.QDiffusionConfig(num_diffusion_timesteps: int = 500, use_coupled_sampling: bool = False, num_candidates: int = 1, proposal_temperature: float = 0.0, proposal_noise_scale: float = 1.0, energy_temperature: float = 1.0, disable_resample: bool = False, resample_ratio: float = 0.25, resample_top_p: float = 0.95, decoding_strategy: str = 'reparam-uncond-deterministic-linear')

Bases: object

Configuration for energy-guided discrete generation.

Attributes:

num_diffusion_timesteps: Number of discrete noising steps used by the

training objective.

use_coupled_sampling: Whether to use the coupled corruption variant.

num_candidates: Number of proposal candidates sampled at each decode step.

proposal_temperature: Temperature used for proposal-side sampling.

proposal_noise_scale: Gumbel noise scale used during proposal sampling.

energy_temperature: Temperature used when converting energies into

reranking weights.

disable_resample: Whether to disable repetition-collapse resampling.

resample_ratio: Frequency threshold that triggers resampling.

resample_top_p: Top-p cutoff used during resampling.

decoding_strategy: Skeptical-remasking strategy string.

decoding_strategy: str = 'reparam-uncond-deterministic-linear'

disable_resample: bool = False

energy_temperature: float = 1.0

num_candidates: int = 1

num_diffusion_timesteps: int = 500

proposal_noise_scale: float = 1.0

proposal_temperature: float = 0.0

resample_ratio: float = 0.25

resample_top_p: float = 0.95

use_coupled_sampling: bool = False

class kaiwu.torch_plugin.QVAE(input_dimension=None, activation_fct=None, config=None, **kwargs)

Bases: AutoEncoderBase

Quantum Variational Autoencoder integrated into AutoEncoderBase framework.

Args:

input_dimension (int or list of int): Dimensionality of input features. activation_fct (callable, optional): Activation function for hidden layers. config (object): Configuration object containing:

• num_latent_units (int)

• loss_type (str): ‘bernoulli’ or ‘mse’

• dist_beta (float, default=10.0)

• kl_beta (float, default=1e-6)

• weight_decay (float, default=0.01)

sampler_type (str): Type of sampler for BM (‘sa’ or ‘cim’). n_batches (int): Number of batches for conditional decoding (0 = no conditioning). bm (object, optional): Pre-created Boltzmann Machine. encoder (object, optional): Pre-created encoder. decoder (object, optional): Pre-created decoder. sampler (object, optional): Pre-created sampler. **kwargs: Additional keyword arguments for AutoEncoderBase.

bm_loss(q, bm_weight_decay=0.0)

Compute BM loss for updating BM parameters separately.

Args:

q (torch.Tensor): Encoder output logits (batch_size, latent_dim).

Must be detached to prevent gradients flowing to encoder.

bm_weight_decay (float, optional): L2 regularization coefficient for BM parameters.

Returns:

torch.Tensor: BM loss scalar.

create_networks()

Create encoder, decoder, BM and sampler. Subclasses must override.

energy(x, loss_type)

计算 BM 能量（用于评估）

forward(x, batch_idx=None)

Forward pass through the QVAE.

Args:

x (torch.Tensor): Input tensor of shape (batch_size, input_dim). batch_idx (torch.Tensor, optional): Batch indices for conditional decoding.

Required if n_batches > 0.

Returns:

tuple: (recon_x, posterior, q, zeta)

• recon_x: Reconstructed logits (batch_size, input_dim)

• posterior: Posterior distribution object (MixtureGeneric)

• q: Encoder output logits (batch_size, latent_dim)

• zeta: Reparameterized latent sample (batch_size, latent_dim)

loss(x, recon_x, posterior)

Compute total loss (reconstruction + KL + weight decay).

Args:

x (torch.Tensor): Input tensor (batch_size, input_dim). recon_x (torch.Tensor): Reconstructed logits (batch_size, input_dim). posterior (MixtureGeneric): Posterior distribution object. q (torch.Tensor): Encoder logits (batch_size, latent_dim). zeta (torch.Tensor): Latent sample (batch_size, latent_dim).

Returns:

torch.Tensor: Total loss scalar.

Raises:

ValueError: If loss_type is not supported.

posterior(q_logits, beta)

Compute posterior distribution and reparameterized sample.

Args:

q_logits (torch.Tensor): Encoder output logits (batch_size, latent_dim). beta (float): Mixture parameter for MixtureGeneric.

Returns:

tuple: (posterior_dist, zeta)

• posterior_dist: MixtureGeneric object

• zeta: Reparameterized sample (batch_size, latent_dim)

set_train_bias(mean)

Compute train bias from dataset mean for Bernoulli reconstruction.

class kaiwu.torch_plugin.RestrictedBoltzmannMachine(num_visible: int, num_hidden: int, quadratic_coef: FloatTensor | None = None, linear_bias: FloatTensor | None = None, device=None)

Bases: AbstractBoltzmannMachine

Create a Restricted Boltzmann Machine.

Args:

num_visible (int): Number of visible nodes in the model.

num_hidden (int): Number of hidden nodes in the model.

quadratic_coef (torch.FloatTensor, optional): quadratic coefficent,

shape is [num_visible, num_hidden]

linear_bias (torch.FloatTensor, optional): linear bias, shape is [num_hidden]

device (torch.device, optional): Device to construct tensors.

clip_parameters(h_range, j_range) → None

Clip linear and quadratic bias weights in-place.

Args:

h_range (tuple[float, float]): Range for quadratic weights. for example, [-1, 1] j_range (tuple[float, float]): Range for linear weights. for example, [-1, 1]

forward(s_all: Tensor) → Tensor

Compute the Hamiltonian.

Args:

s_all (torch.tensor): Tensor of shape (B, N), where B is the batch size,

and N is the number of variables in the model.

Returns:

torch.tensor: Hamiltonian of shape (B,).

get_hidden(s_visible: Tensor, requires_grad: bool = False, bernoulli: bool = False) → Tensor

Propagate visible spins to the hidden layer.

Args:

s_visible: Visible layer tensor. requires_grad: Whether to allow gradient backpropagation.

get_visible(s_hidden: Tensor, bernoulli: bool = False) → Tensor

Propagate hidden spins to the visible layer.

property hidden_bias: Tensor

Return the hidden bias.

property visible_bias: Tensor

Return the visible bias.

class kaiwu.torch_plugin.UnsupervisedDBN(hidden_layers_structure=None)

Bases: Module

A general unsupervised Deep Belief Network (DBN) architecture.

This model is a stack of Restricted Boltzmann Machines (RBMs).

Args:

hidden_layers_structure (list, optional): A list of integers

representing the number of hidden units in each layer. Defaults to [100, 100].

create_rbm_layer(input_dim)

Creates the layers of RBMs for the DBN.

Args:

input_dim (int): The dimension of the input data (number of visible units).

Returns:

UnsupervisedDBN: The instance itself with the RBM layers created.

forward(data_in)

Performs a forward pass to transform the input data.

Args:

data_in (numpy.ndarray): The input data.

Returns:

numpy.ndarray: The transformed data after passing through all RBM layers.

Raises:

ValueError: If the model has not been built or trained yet.

get_rbm_layer(index)

Gets the RBM layer at the specified index.

Args:

index (int): The index of the RBM layer.

Returns:

RestrictedBoltzmannMachine or None: The RBM layer if found, otherwise None.

mark_as_trained()

Marks the model as trained.

Returns:

UnsupervisedDBN: The instance itself.

property num_layers

Returns the number of RBM layers.

Returns:

int: The number of layers.

property output_dim

Returns the output dimension of the DBN.

Returns:

int: The dimension of the final hidden layer.

reconstruct(data_in, layer_index=0)

Reconstructs the input from a specified RBM layer.

Args:

data_in (numpy.ndarray): The input data to be reconstructed.

layer_index (int, optional): The index of the RBM layer to use for reconstruction.

Defaults to 0.

Returns:

numpy.ndarray: The reconstructed data.

Raises:

ValueError: If the model has no RBM layers or the layer index is out of range.

static reconstruct_with_rbm(rbm, data_in, device=None)

Reconstructs data using a single RBM.

Args:

rbm (RestrictedBoltzmannMachine): The trained RBM model.

data_in (numpy.ndarray): The input data.

device (torch.device, optional): The device to perform computation on.

If None, uses the RBM’s device. Defaults to None.

Returns:

tuple[numpy.ndarray, numpy.ndarray]: A tuple containing:

• The reconstructed visible layer data.

• The reconstruction error for each sample.

transform(data_in)

An sklearn-compatible transform method.

Args:

data_in (numpy.ndarray): The input data.

Returns:

numpy.ndarray: The transformed data.

kaiwu.torch_plugin.disable_usage_stats() → None

Disable kpp usage statistics. No data will be reported after disabling.

kaiwu.torch_plugin.enable_usage_stats() → None

Enable kpp usage statistics (enabled by default; usually not needed).

kaiwu.torch_plugin.is_usage_stats_enabled() → bool

Return whether usage statistics are currently enabled.