use trustformers_core::{
errors::Result,
layers::{LayerNorm, Linear},
quantum::{QuantumAnsatz, QuantumNeuralLayer},
tensor::Tensor,
traits::Layer,
};
use super::{config::QuantumClassicalConfig, model::QuantumClassicalModelOutput};
#[derive(Debug)]
pub struct QuantumConvolutionalNN {
pub config: QuantumClassicalConfig,
pub conv_layers: Vec<Linear>,
pub quantum_pooling_layers: Vec<QuantumNeuralLayer>,
pub fc_layers: Vec<Linear>,
pub layer_norms: Vec<LayerNorm>,
pub output_projection: Linear,
}
impl QuantumConvolutionalNN {
pub fn new(config: &QuantumClassicalConfig) -> Result<Self> {
let mut conv_layers = Vec::new();
let mut quantum_pooling_layers = Vec::new();
let mut fc_layers = Vec::new();
let mut layer_norms = Vec::new();
for _ in 0..config.n_classical_layers {
conv_layers.push(Linear::new(config.d_model, config.d_model, config.use_bias));
fc_layers.push(Linear::new(config.d_model, config.d_model, config.use_bias));
layer_norms.push(LayerNorm::new(vec![config.d_model], 1e-12)?);
}
for _ in 0..config.n_quantum_layers {
let ansatz = QuantumAnsatz::from(config.quantum_ansatz.clone());
let parameters = vec![0.1; config.get_quantum_parameters_count()];
quantum_pooling_layers.push(QuantumNeuralLayer::new(
config.num_qubits,
ansatz,
¶meters,
)?);
}
let output_projection = Linear::new(config.d_model, config.d_model, config.use_bias);
Ok(Self {
config: config.clone(),
conv_layers,
quantum_pooling_layers,
fc_layers,
layer_norms,
output_projection,
})
}
pub fn forward(&mut self, input: &Tensor) -> Result<QuantumClassicalModelOutput> {
let mut hidden_states = input.clone();
let mut quantum_measurements = Vec::new();
for (i, (conv_layer, layer_norm)) in
self.conv_layers.iter().zip(self.layer_norms.iter()).enumerate()
{
let conv_output = conv_layer.forward(hidden_states.clone())?;
hidden_states = layer_norm.forward(conv_output)?;
if i < self.quantum_pooling_layers.len() {
let quantum_output = self.quantum_pooling_layers[i].forward(&hidden_states)?;
quantum_measurements.push(quantum_output.clone());
hidden_states = quantum_output;
}
}
for fc_layer in &self.fc_layers {
hidden_states = fc_layer.forward(hidden_states)?;
hidden_states = hidden_states.relu()?;
}
let output = self.output_projection.forward(hidden_states.clone())?;
let combined_measurements = if !quantum_measurements.is_empty() {
let mut combined = quantum_measurements[0].clone();
for measurement in quantum_measurements.iter().skip(1) {
combined = combined.add(measurement)?;
}
Some(combined)
} else {
None
};
Ok(QuantumClassicalModelOutput {
hidden_states: output,
quantum_measurements: combined_measurements,
classical_activations: Some(hidden_states),
quantum_attention_weights: None,
quantum_fidelity_scores: None,
quantum_entanglement_measures: None,
quantum_error_mitigation: None,
})
}
pub fn parameter_count(&self) -> usize {
self.conv_layers.iter().map(|l| l.parameter_count()).sum::<usize>()
+ self.fc_layers.iter().map(|l| l.parameter_count()).sum::<usize>()
+ self.layer_norms.iter().map(|l| l.parameter_count()).sum::<usize>()
+ self.output_projection.parameter_count()
+ self.quantum_pooling_layers.len() * self.config.get_quantum_parameters_count()
}
pub fn memory_usage(&self) -> f32 {
self.parameter_count() as f32 * 4.0 / 1_000_000.0
}
}