API Reference

Models

class models.product.BinaryProductModel(n_inputs, weight_constraint=None, use_gaussian_init=True, gaussian_mean=0.5, gaussian_std=0.25)

Bases: Module

Single-output binary product model (DEPRECATED - use MultiBinaryProductModel with n_outputs=1).

Parameters:

• n_inputs – Number of inputs

• weight_constraint – Optional constraint function for weights

• use_gaussian_init – If True, use Gaussian initialization instead of Uniform[0,1]

• gaussian_mean – Mean for Gaussian initialization (default: 0.5)

• gaussian_std – Standard deviation for Gaussian initialization (default: 0.25)

forward(inputs)

Forward pass.

class models.product.MultiBinaryProductModel(n_outputs, weight_constraint=None, use_gaussian_init=True, gaussian_mean=0.5, gaussian_std=0.25)

Bases: Module

Multi-output binary product model for parallel XOR computation.

Parameters:

• n_outputs – Number of parallel XOR units

• weight_constraint – Optional constraint function for weights

• use_gaussian_init – If True, use Gaussian initialization instead of Uniform[0,1]

• gaussian_mean – Mean for Gaussian initialization (default: 0.5)

• gaussian_std – Standard deviation for Gaussian initialization (default: 0.25)

forward(inputs)

Forward pass.

Parameters:

inputs – Tensor of shape [batch_size, n_inputs] with binary values # THATS THE THEORY BUT NOTHING ENFORCES IT IN THE MODEL

Returns:

XOR outputs of shape [batch_size, n_outputs]

class models.product.MultiBinaryProductModelWithOracle(n_outputs, p_e=0.0, weight_constraint=None, use_gaussian_init=True, gaussian_mean=0.5, gaussian_std=0.25)

Bases: Module

Multi-output binary product model with oracle for supervised XOR learning.

This model includes: - A BSC channel to add noise to inputs - An oracle (non-trainable) that computes true XOR outputs - A trainable model that learns to replicate the oracle

Parameters:

• n_outputs – Number of parallel XOR units

• p_e – Error probability for the BSC channel (default: 0.0)

• weight_constraint – Optional constraint function for weights

• use_gaussian_init – If True, use Gaussian initialization instead of Uniform[0,1]

• gaussian_mean – Mean for Gaussian initialization (default: 0.5)

• gaussian_std – Standard deviation for Gaussian initialization (default: 0.25)

forward(inputs)

Forward pass.

Parameters:

inputs – Tensor of shape [batch_size, n_inputs] with binary values

Returns:

Tuple of (xor_oracle, xor_model, p_epsilon) - xor_oracle: Oracle outputs [batch_size, n_outputs] - xor_model: Model outputs [batch_size, n_outputs] - p_epsilon: Parameter error rate per XOR unit [n_outputs]

set_bsc_error_probability(p_e)

Set the BSC error probability.

set_model_parameters(model_weights)

Set the model weights (trainable).

set_oracle_parameters(oracle_weights)

Set the oracle weights (ground truth).

class models.mlp.ParallelMLP(n_inputs, n_outputs, hidden_sizes=(512, 512, 64, 1), activation=None, output_activation=None, use_bias=True)

Bases: Module

Parallel Multi-Layer Perceptron with independent networks computed efficiently.

Instead of running n_outputs separate MLPs sequentially, this implementation processes all outputs in parallel using vectorized operations.

Parameters:

• n_inputs – Number of input features

• n_outputs – Number of parallel independent networks

• hidden_sizes – Tuple of hidden layer sizes (e.g., (512, 512, 64))

• activation – Activation function to use (default: nn.ReLU())

• output_activation – Optional activation for output layer (default: None)

• use_bias – Whether to use bias terms (default: True)

Example

>>> model = ParallelMLP(n_inputs=10, n_outputs=4, hidden_sizes=(64, 32))
>>> inputs = torch.randn(5, 10)  # batch_size=5
>>> outputs = model(inputs)  # shape: [5, 4, 1]

forward(inputs)

Forward pass through parallel networks.

Parameters:

inputs – Tensor of shape [batch_size, n_inputs]

Returns:

Tensor of shape [batch_size, n_outputs, 1] (or [batch_size, n_outputs] if squeezed)

set_weights(layer_idx, weights, biases=None)

Set weights for a specific layer.

Parameters:

• layer_idx – Index of the layer (0-indexed)

• weights – Tensor of shape [n_outputs, out_features, in_features]

• biases – Optional tensor of shape [n_outputs, out_features]

class models.layers.product.BinaryProductLayer(n_inputs, weight_constraint=None, use_gaussian_init=True, gaussian_mean=0.5, gaussian_std=0.25)

Bases: Module

Single-output binary product layer (DEPRECATED - use MultiBinaryProductLayer with n_outputs=1).

Parameters:

• n_inputs – Number of inputs

• weight_constraint – Optional constraint function for weights

• use_gaussian_init – If True, use Gaussian initialization instead of Uniform[0,1]

• gaussian_mean – Mean for Gaussian initialization (default: 0.5)

• gaussian_std – Standard deviation for Gaussian initialization (default: 0.25)

forward(inputs)

Forward pass.

Parameters:

inputs – Tensor of shape [batch_size, n_inputs] with binary values {0, 1} # THATS THE THEORY BUT NOTHING ENFORCES IT IN THE MODEL

Returns:

Tensor of shape [batch_size] with values in [0, 1] # THATS THE THEORY BUT NOTHING ENFORCES IT IN THE MODEL

set_xor_parameters(xor_weights)

Set weights directly.

Parameters:

xor_weights – Tensor of shape [n_inputs] with XOR weights

class models.layers.product.MultiBinaryProductLayer(n_outputs, weight_constraint=None, hard_step=False, use_gaussian_init=True, gaussian_mean=0.5, gaussian_std=0.25)

Bases: Module

Multi-output binary product layer for XOR computation.

Implements the product node: y = prod_i (w_i * z_i + 1) where z_i = x_i - 1 This provides a continuous extension of XOR operation.

Parameters:

• n_outputs – Number of parallel XOR units

• weight_constraint – Optional constraint function for weights

• hard_step – If True, use hard threshold at 0.5 (non-differentiable)

• use_gaussian_init – If True, use Gaussian initialization instead of Uniform[0,1]

• gaussian_mean – Mean for Gaussian initialization (default: 0.5)

• gaussian_std – Standard deviation for Gaussian initialization (default: 0.25)

forward(inputs)

Forward pass of the product layer.

Parameters:

inputs – Tensor of shape [batch_size, n_inputs] with binary values {0, 1} # THATS THE THEORY BUT NOTHING ENFORCES IT IN THE MODEL

Returns:

Tensor of shape [batch_size, n_outputs] with values in [0, 1] # THATS THE THEORY BUT NOTHING ENFORCES IT IN THE MODEL

set_xor_parameters(xor_weights)

Set weights to specific values.

Parameters:

xor_weights – Tensor of shape [n_inputs, n_outputs] with XOR weights

class models.layers.channels.BinarySymmetricChannelLayer(p_e)

Bases: Module

Binary Symmetric Channel layer that flips bits with probability p_e.

This simulates a BSC channel where each bit is flipped independently with probability p_e.

Parameters:

p_e – Error probability (probability of bit flip)

forward(inputs)

Apply BSC noise to inputs.

Parameters:

inputs – Tensor of binary values {0, 1} # AGAIN NOTHING ENFORCE IT IN THE MODEL

Returns:

Noisy tensor with bits flipped according to probability p_e # AGAIN NOTHING ENFORCE IT IN THE MODEL

set_error_probability(p_e)

Update the error probability.

Training

training.training.create_and_train_model(n_inputs, n_outputs, p_e, learning_rate=0.1, batch_size=100, max_steps=1000, convergence_threshold=0.01, stagnation_window=-1, stagnation_threshold=1e-06, device='mps', seed=42, verbose=True, print_interval=100, record_history=False, record_weights=False, p_w=0.5, use_gaussian_init=True, gaussian_mean=0.5, gaussian_std=0.25, match_oracle_weights=False)

Create, initialize, and train a model with the specified parameters.

Parameters:

• n_inputs – Number of input features

• n_outputs – Number of parallel XOR units

• oracle_weights – Oracle weights tensor [n_inputs, n_outputs]

• inputs – Input data tensor [batch_size, n_inputs]

• learning_rate – Learning rate for optimizer

• max_steps – Maximum number of training steps

• convergence_threshold – Convergence threshold for mean p_diff

• p_e – Error probability for BSC channel (default: 0.0)

• stagnation_window – Window size for stagnation detection (default: -1 = disabled)

• stagnation_threshold – Minimum improvement per step (default: 1e-3)

• device – Device to use (‘cpu’, ‘cuda’, ‘mps’) (default: ‘cpu’)

• seed – Random seed for reproducibility (default: 42)

• verbose – Whether to print progress (default: True)

• print_interval – Print progress every N steps (default: 100)

• record_history – Whether to record full training history (default: False)

• record_weights – Whether to record weights during training (default: False)

• p_w – proportion of Oracle weights equal to 1 (default: 0.5),

• use_gaussian_init – Whether to use Gaussian initialization for model weights (default: True)

• gaussian_mean – Mean of Gaussian initialization (default: 0.5)

• gaussian_std – Standard deviation of Gaussian initialization (default: 0.25)

• match_oracle_weights – Whether to initialize model weights to match oracle weights with dispersion following gaussian init. If gaussian mean and std equal 0, then exact match (default: False)

Returns:

• training_results: Results from train_until_convergence

Return type:

Dictionary containing

training.training.train_until_convergence(model, optimizer, inputs, max_steps, convergence_threshold, stagnation_window=-1, stagnation_threshold=0.001, verbose=True, print_interval=100, record_history=False, record_weights=False)

Train model until p_diff < convergence_threshold or max_steps reached.

Parameters:

• model – Model to train (must return y_oracle, y_model, p_epsilon, p_diff)

• optimizer – Optimizer instance

• inputs – Input data tensor [batch_size, n_inputs]

• max_steps – Maximum number of training steps

• convergence_threshold – Convergence threshold for mean p_diff

• stagnation_window – Window size for checking improvement rate (default: -1 = disabled) If > 0, training stops early if mean improvement over window < stagnation_threshold

• stagnation_threshold – Minimum mean improvement per step required (default: 1e-3) Rate is measured as: (p_diff[start] - p_diff[end]) / window_size

• verbose – Whether to print progress (default: True)

• print_interval – Print progress every N steps (default: 100)

• record_history – Whether to record full training history (default: False)

• record_weights – Whether to record weight during training (default: False)

Returns:

• steps: Number of steps taken (max_steps if didn’t converge or stagnated)

• final_p_diff: Final mean p_diff value

• final_p_epsilon: Final mean p_epsilon value

• final_loss: Final loss value

• converged: Whether convergence was reached

• stagnated: Whether stagnation was detected

• oracle_weights: Oracle weights tensor [n_inputs, n_outputs] (if record_weights=True)

• model_weights_init: Model weights tensor at init [n_inputs, n_outputs] (if record_weights=True)

• history: Training history (if record_history=True), contains:

  • p_diff: List of mean p_diff values per step

  • p_epsilon: List of mean p_epsilon values per step

  • loss: List of loss values per step

  • model_weights: List of model weights at recorded steps (if record_weights=True)

Return type:

Dictionary containing

Plotting

Plotting utilities

plotting.plot_utility.compute_kde_envelope(weights, x_range)

Compute KDE envelope for given weights.

plotting.plot_utility.compute_steps_to_thresholds(x_param_list, results, n_thresholds=100)

Compute steps needed to reach different metric thresholds (preprocessing for inverted plots).

This function inverts the data: instead of metric values at each step, it computes the step number needed to reach each threshold.

Parameters:

• x_param_list – List of x parameter values (e.g., p_e values)

• results – 2D array [n_params, n_steps] with metric values (decreasing from ~1 to ~0)

• n_thresholds – Number of metric thresholds to compute

Returns:

Array of threshold values result_steps: 2D array [n_params, n_thresholds] with steps to reach each threshold

Return type:

thresholds

plotting.plot_utility.create_subplot_grid(rows, cols)

Create subplot grid using default figsize scaled by rows/cols.

plotting.plot_utility.enable_minor_ticks()

Manually enable minor ticks on current axes (if hook doesn’t work).

plotting.plot_utility.extract_data(results_dict, metric='p_diff')

Extract data from training results dictionary.

Parameters:

• results_dict – Dict {x_param: {‘training_results’: {‘history’: {…}}}}

• metric – Metric to extract (‘p_diff’, ‘p_epsilon’, ‘loss’)

Returns:

List of x parameter values (sorted keys from dict) results: 2D array [n_params, max_steps] (padded with None)

Return type:

x_param_list

plotting.plot_utility.plot_distribution_analysis(snapshots, n_oracle_0, n_oracle_1, text_ratio=1.2)

plotting.plot_utility.plot_metric_vs_param(x_param_list, results, ax, draw_best=True, best='min', xscale='log', yscale='linear', xlim=None, ylim=None, value_labels=None, add_min_curve=False)

Plot metric vs parameter at different steps (iso-step curves) or thresholds (iso-metric curves).

Generic plotting function that reproduces original subplot() function exactly.

Parameters:

• x_param_list – List of x parameter values (e.g., p_e values)

• results – 2D array [n_params, n_steps_or_thresholds] with metric values or steps

• ax – Matplotlib axis to plot on

• draw_best – Whether to plot best reached metric (default: True)

• best – ‘min’ or ‘max’ to indicate whether best is minimum or maximum (default: ‘min’)

• xscale – X-axis scale (default: ‘log’)

• yscale – Y-axis scale (default: ‘linear’)

• xlim – tuple, optional X-axis limits as (min, max)

• ylim – tuple, optional Y-axis limits as (min, max)

• value_labels – Optional list of labels for each curve (e.g., threshold values for inverted plots)

• add_min_curve – Whether to add min curve curve (for inverted plots)

plotting.plot_utility.plot_qq_with_r2_list(ax, snapshots, colors, weight_key, title, n_sample=500, text_ratio=1.2)

Plot Q-Q with undersampling for visibility and full R² computation.

plotting.plot_utility.reset_plot_style()

Reset to matplotlib defaults.

plotting.plot_utility.set_figure_size(width=10, height=6)

Set figure size for current or next figure.

plotting.plot_utility.setup_plot_style()

Setup comprehensive matplotlib style with minor ticks enabled by default.