osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
792 lines
24 KiB
Rust
792 lines
24 KiB
Rust
//! Module composition strategies for building complex architectures
|
|
//!
|
|
//! This module implements various ways to compose modules together,
|
|
//! enabling the creation of complex computational graphs from simple building blocks.
|
|
|
|
use super::*;
|
|
|
|
/// Strategy for composing multiple modules
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum CompositionStrategy {
|
|
/// Sequential composition (pipeline)
|
|
Sequential,
|
|
|
|
/// Parallel composition with combination
|
|
Parallel(CombinationStrategy),
|
|
|
|
/// Hierarchical composition (multi-level)
|
|
Hierarchical,
|
|
|
|
/// Graph-based composition with arbitrary connections
|
|
Graph,
|
|
|
|
/// Dynamic composition based on input
|
|
Dynamic,
|
|
}
|
|
|
|
/// Strategy for combining outputs from parallel modules
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum CombinationStrategy {
|
|
/// Concatenate outputs along feature dimension
|
|
Concatenate,
|
|
|
|
/// Element-wise average
|
|
Average,
|
|
|
|
/// Element-wise sum
|
|
Sum,
|
|
|
|
/// Weighted sum with fixed weights
|
|
WeightedSum(Vec<f32>),
|
|
|
|
/// Attention-weighted combination
|
|
AttentionWeighted,
|
|
|
|
/// Maximum pooling across modules
|
|
Max,
|
|
|
|
/// Learned combination using a neural network
|
|
Learned,
|
|
}
|
|
|
|
impl CombinationStrategy {
|
|
pub fn combine(&self, tensors: &[Tensor]) -> Result<Tensor> {
|
|
if tensors.is_empty() {
|
|
return Err(ModularError::CompositionError {
|
|
message: "No tensors to combine".to_string(),
|
|
});
|
|
}
|
|
|
|
if tensors.len() == 1 {
|
|
return Ok(tensors[0].clone());
|
|
}
|
|
|
|
match self {
|
|
CombinationStrategy::Concatenate => {
|
|
// Concatenate along the last dimension (feature dimension)
|
|
Tensor::cat(tensors, -1).map_err(|e| ModularError::CompositionError {
|
|
message: format!("Concatenation failed: {}", e),
|
|
})
|
|
}
|
|
|
|
CombinationStrategy::Average => {
|
|
let sum = tensors
|
|
.iter()
|
|
.try_fold(tensors[0].clone(), |acc, t| acc.add(t))?;
|
|
sum.div_scalar(tensors.len() as f32)
|
|
.map_err(|e| ModularError::CompositionError {
|
|
message: format!("Average failed: {}", e),
|
|
})
|
|
}
|
|
|
|
CombinationStrategy::Sum => tensors
|
|
.iter()
|
|
.try_fold(tensors[0].clone(), |acc, t| acc.add(t))
|
|
.map_err(|e| ModularError::CompositionError {
|
|
message: format!("Sum failed: {}", e),
|
|
}),
|
|
|
|
CombinationStrategy::WeightedSum(weights) => {
|
|
if weights.len() != tensors.len() {
|
|
return Err(ModularError::CompositionError {
|
|
message: format!(
|
|
"Weight count {} doesn't match tensor count {}",
|
|
weights.len(),
|
|
tensors.len()
|
|
),
|
|
});
|
|
}
|
|
|
|
let mut result = tensors[0].mul_scalar(weights[0])?;
|
|
for (i, tensor) in tensors.iter().enumerate().skip(1) {
|
|
let weighted = tensor.mul_scalar(weights[i])?;
|
|
result = result.add(&weighted)?;
|
|
}
|
|
Ok(result)
|
|
}
|
|
|
|
CombinationStrategy::Max => {
|
|
// Stack tensors and take max along the new dimension
|
|
let stacked = Tensor::stack(tensors, 0)?;
|
|
stacked.max().map_err(|e| ModularError::CompositionError {
|
|
message: format!("Max pooling failed: {}", e),
|
|
})
|
|
}
|
|
|
|
CombinationStrategy::AttentionWeighted => {
|
|
// Simplified attention-based combination
|
|
// In practice, this would use a proper attention mechanism
|
|
self.combine(&[tensors[0].clone()]) // Fallback to first tensor
|
|
}
|
|
|
|
CombinationStrategy::Learned => {
|
|
// Placeholder - would use a learned combination network
|
|
self.combine(&[tensors[0].clone()]) // Fallback to first tensor
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Result of module composition
|
|
#[derive(Debug)]
|
|
pub struct CompositionResult {
|
|
/// Final output tensor
|
|
pub output: Tensor,
|
|
|
|
/// Intermediate outputs from each stage/module
|
|
pub intermediate_outputs: Vec<Tensor>,
|
|
|
|
/// Modules that were actually used
|
|
pub active_modules: Vec<String>,
|
|
|
|
/// Total computational cost
|
|
pub total_cost: f32,
|
|
|
|
/// Composition metadata
|
|
pub metadata: std::collections::HashMap<String, String>,
|
|
}
|
|
|
|
/// Configuration for module composition
|
|
#[derive(Debug, Clone)]
|
|
pub struct CompositionConfig {
|
|
/// Whether to store intermediate outputs
|
|
pub store_intermediates: bool,
|
|
|
|
/// Maximum depth for hierarchical composition
|
|
pub max_depth: usize,
|
|
|
|
/// Whether to enable gradient flow through all paths
|
|
pub enable_full_gradients: bool,
|
|
|
|
/// Resource limits for composition
|
|
pub resource_limits: ResourceBudget,
|
|
}
|
|
|
|
impl Default for CompositionConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
store_intermediates: false,
|
|
max_depth: 5,
|
|
enable_full_gradients: true,
|
|
resource_limits: ResourceBudget::default(),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Sequential Composition
|
|
// ============================================================================
|
|
|
|
/// Sequential composition of modules (pipeline)
|
|
#[derive(Debug)]
|
|
pub struct SequentialComposition {
|
|
modules: Vec<Box<dyn Module>>,
|
|
config: CompositionConfig,
|
|
metadata: ModuleMetadata,
|
|
}
|
|
|
|
impl SequentialComposition {
|
|
pub fn new(modules: Vec<Box<dyn Module>>) -> Self {
|
|
Self {
|
|
modules,
|
|
config: CompositionConfig::default(),
|
|
metadata: ModuleMetadata::default(),
|
|
}
|
|
}
|
|
|
|
pub fn with_config(mut self, config: CompositionConfig) -> Self {
|
|
self.config = config;
|
|
self
|
|
}
|
|
|
|
pub fn forward(&self, input: &Tensor) -> Result<CompositionResult> {
|
|
if self.modules.is_empty() {
|
|
return Err(ModularError::CompositionError {
|
|
message: "No modules in sequential composition".to_string(),
|
|
});
|
|
}
|
|
|
|
let mut current_input = input.clone();
|
|
let mut intermediate_outputs = Vec::new();
|
|
let mut active_modules = Vec::new();
|
|
let mut total_cost = 0.0;
|
|
|
|
for module in &self.modules {
|
|
// Check compatibility
|
|
if current_input.dims().last() != Some(&module.input_dim()) {
|
|
return Err(ModularError::IncompatibleDimensions {
|
|
expected: module.input_dim(),
|
|
actual: *current_input.dims().last().unwrap(),
|
|
});
|
|
}
|
|
|
|
// Forward pass through module
|
|
let output = module.forward(¤t_input)?;
|
|
|
|
// Update state
|
|
if self.config.store_intermediates {
|
|
intermediate_outputs.push(current_input.clone());
|
|
}
|
|
active_modules.push(module.module_id().to_string());
|
|
total_cost += module.computation_cost(current_input.dims());
|
|
|
|
current_input = output;
|
|
}
|
|
|
|
let mut metadata = std::collections::HashMap::new();
|
|
metadata.insert("composition_type".to_string(), "sequential".to_string());
|
|
metadata.insert("num_modules".to_string(), self.modules.len().to_string());
|
|
|
|
Ok(CompositionResult {
|
|
output: current_input,
|
|
intermediate_outputs,
|
|
active_modules,
|
|
total_cost,
|
|
metadata,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl Module for SequentialComposition {
|
|
fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
|
Ok(self.forward(input)?.output)
|
|
}
|
|
|
|
fn input_dim(&self) -> usize {
|
|
self.modules.first().map(|m| m.input_dim()).unwrap_or(0)
|
|
}
|
|
|
|
fn output_dim(&self) -> usize {
|
|
self.modules.last().map(|m| m.output_dim()).unwrap_or(0)
|
|
}
|
|
|
|
fn module_type(&self) -> &str {
|
|
"sequential_composition"
|
|
}
|
|
|
|
fn module_id(&self) -> &str {
|
|
"sequential"
|
|
}
|
|
|
|
fn capabilities(&self) -> Vec<String> {
|
|
let mut all_caps = Vec::new();
|
|
for module in &self.modules {
|
|
all_caps.extend(module.capabilities());
|
|
}
|
|
all_caps.sort();
|
|
all_caps.dedup();
|
|
all_caps
|
|
}
|
|
|
|
fn metadata(&self) -> &ModuleMetadata {
|
|
self.modules
|
|
.first()
|
|
.map(|m| m.metadata())
|
|
.unwrap_or(&self.metadata)
|
|
}
|
|
|
|
fn complexity_score(&self) -> f32 {
|
|
self.modules.iter().map(|m| m.complexity_score()).sum()
|
|
}
|
|
|
|
fn memory_footprint(&self) -> usize {
|
|
self.modules.iter().map(|m| m.memory_footprint()).sum()
|
|
}
|
|
|
|
fn is_stateful(&self) -> bool {
|
|
self.modules.iter().any(|m| m.is_stateful())
|
|
}
|
|
|
|
fn requires_training(&self) -> bool {
|
|
self.modules.iter().any(|m| m.requires_training())
|
|
}
|
|
|
|
fn parameters(&self) -> Vec<&Tensor> {
|
|
self.modules.iter().flat_map(|m| m.parameters()).collect()
|
|
}
|
|
|
|
fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
|
|
self.modules
|
|
.iter_mut()
|
|
.flat_map(|m| m.parameters_mut())
|
|
.collect()
|
|
}
|
|
|
|
fn clone_module(&self) -> Box<dyn Module> {
|
|
let cloned_modules = self.modules.iter().map(|m| m.clone_module()).collect();
|
|
|
|
Box::new(SequentialComposition {
|
|
modules: cloned_modules,
|
|
config: self.config.clone(),
|
|
metadata: self.metadata.clone(),
|
|
})
|
|
}
|
|
|
|
fn specialize_for_task(
|
|
&self,
|
|
task_type: TaskType,
|
|
training_data: &[TestData],
|
|
) -> Result<Box<dyn Module>> {
|
|
let specialized_modules = self
|
|
.modules
|
|
.iter()
|
|
.map(|m| m.specialize_for_task(task_type.clone(), training_data))
|
|
.collect::<Result<Vec<_>>>()?;
|
|
|
|
Ok(Box::new(SequentialComposition::new(specialized_modules)))
|
|
}
|
|
|
|
fn continually_specialize(&mut self, training_data: &[TestData], task_id: usize) -> Result<()> {
|
|
for module in &mut self.modules {
|
|
module.continually_specialize(training_data, task_id)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Parallel Composition
|
|
// ============================================================================
|
|
|
|
/// Parallel composition of modules with combination strategy
|
|
#[derive(Debug)]
|
|
pub struct ParallelComposition {
|
|
modules: Vec<Box<dyn Module>>,
|
|
combination_strategy: CombinationStrategy,
|
|
config: CompositionConfig,
|
|
metadata: ModuleMetadata,
|
|
}
|
|
|
|
impl ParallelComposition {
|
|
pub fn new(modules: Vec<Box<dyn Module>>, combination_strategy: CombinationStrategy) -> Self {
|
|
Self {
|
|
modules,
|
|
combination_strategy,
|
|
config: CompositionConfig::default(),
|
|
metadata: ModuleMetadata::default(),
|
|
}
|
|
}
|
|
|
|
pub fn with_config(mut self, config: CompositionConfig) -> Self {
|
|
self.config = config;
|
|
self
|
|
}
|
|
|
|
pub fn forward(&self, input: &Tensor) -> Result<CompositionResult> {
|
|
if self.modules.is_empty() {
|
|
return Err(ModularError::CompositionError {
|
|
message: "No modules in parallel composition".to_string(),
|
|
});
|
|
}
|
|
|
|
let mut outputs = Vec::new();
|
|
let mut active_modules = Vec::new();
|
|
let mut total_cost = 0.0;
|
|
|
|
// Execute all modules in parallel
|
|
for module in &self.modules {
|
|
// Check input compatibility
|
|
if input.dims().last() != Some(&module.input_dim()) {
|
|
return Err(ModularError::IncompatibleDimensions {
|
|
expected: module.input_dim(),
|
|
actual: *input.dims().last().unwrap(),
|
|
});
|
|
}
|
|
|
|
let output = module.forward(input)?;
|
|
outputs.push(output);
|
|
active_modules.push(module.module_id().to_string());
|
|
total_cost += module.computation_cost(input.dims());
|
|
}
|
|
|
|
// Combine outputs
|
|
let combined_output = self.combination_strategy.combine(&outputs)?;
|
|
|
|
let mut metadata = std::collections::HashMap::new();
|
|
metadata.insert("composition_type".to_string(), "parallel".to_string());
|
|
metadata.insert("num_modules".to_string(), self.modules.len().to_string());
|
|
metadata.insert(
|
|
"combination_strategy".to_string(),
|
|
format!("{:?}", self.combination_strategy),
|
|
);
|
|
|
|
Ok(CompositionResult {
|
|
output: combined_output,
|
|
intermediate_outputs: if self.config.store_intermediates {
|
|
outputs
|
|
} else {
|
|
Vec::new()
|
|
},
|
|
active_modules,
|
|
total_cost,
|
|
metadata,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl Module for ParallelComposition {
|
|
fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
|
Ok(self.forward(input)?.output)
|
|
}
|
|
|
|
fn input_dim(&self) -> usize {
|
|
// All modules should have the same input dimension
|
|
self.modules.first().map(|m| m.input_dim()).unwrap_or(0)
|
|
}
|
|
|
|
fn output_dim(&self) -> usize {
|
|
match &self.combination_strategy {
|
|
CombinationStrategy::Concatenate => self.modules.iter().map(|m| m.output_dim()).sum(),
|
|
_ => {
|
|
// For other strategies, output dimension equals input dimension of modules
|
|
self.modules.first().map(|m| m.output_dim()).unwrap_or(0)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn module_type(&self) -> &str {
|
|
"parallel_composition"
|
|
}
|
|
|
|
fn module_id(&self) -> &str {
|
|
"parallel"
|
|
}
|
|
|
|
fn capabilities(&self) -> Vec<String> {
|
|
let mut all_caps = Vec::new();
|
|
for module in &self.modules {
|
|
all_caps.extend(module.capabilities());
|
|
}
|
|
all_caps.sort();
|
|
all_caps.dedup();
|
|
all_caps
|
|
}
|
|
|
|
fn metadata(&self) -> &ModuleMetadata {
|
|
self.modules
|
|
.first()
|
|
.map(|m| m.metadata())
|
|
.unwrap_or(&self.metadata)
|
|
}
|
|
|
|
fn complexity_score(&self) -> f32 {
|
|
self.modules
|
|
.iter()
|
|
.map(|m| m.complexity_score())
|
|
.fold(0.0, f32::max) // Max complexity since they run in parallel
|
|
}
|
|
|
|
fn memory_footprint(&self) -> usize {
|
|
self.modules.iter().map(|m| m.memory_footprint()).sum() // All modules loaded simultaneously
|
|
}
|
|
|
|
fn is_stateful(&self) -> bool {
|
|
self.modules.iter().any(|m| m.is_stateful())
|
|
}
|
|
|
|
fn requires_training(&self) -> bool {
|
|
self.modules.iter().any(|m| m.requires_training())
|
|
}
|
|
|
|
fn parameters(&self) -> Vec<&Tensor> {
|
|
self.modules.iter().flat_map(|m| m.parameters()).collect()
|
|
}
|
|
|
|
fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
|
|
self.modules
|
|
.iter_mut()
|
|
.flat_map(|m| m.parameters_mut())
|
|
.collect()
|
|
}
|
|
|
|
fn clone_module(&self) -> Box<dyn Module> {
|
|
let cloned_modules = self.modules.iter().map(|m| m.clone_module()).collect();
|
|
|
|
Box::new(ParallelComposition {
|
|
modules: cloned_modules,
|
|
combination_strategy: self.combination_strategy.clone(),
|
|
config: self.config.clone(),
|
|
metadata: ModuleMetadata::default(),
|
|
})
|
|
}
|
|
|
|
fn specialize_for_task(
|
|
&self,
|
|
task_type: TaskType,
|
|
training_data: &[TestData],
|
|
) -> Result<Box<dyn Module>> {
|
|
let specialized_modules = self
|
|
.modules
|
|
.iter()
|
|
.map(|m| m.specialize_for_task(task_type.clone(), training_data))
|
|
.collect::<Result<Vec<_>>>()?;
|
|
|
|
Ok(Box::new(ParallelComposition::new(
|
|
specialized_modules,
|
|
self.combination_strategy.clone(),
|
|
)))
|
|
}
|
|
|
|
fn continually_specialize(&mut self, training_data: &[TestData], task_id: usize) -> Result<()> {
|
|
for module in &mut self.modules {
|
|
module.continually_specialize(training_data, task_id)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Hierarchical Composition
|
|
// ============================================================================
|
|
|
|
/// Hierarchical composition with multiple levels
|
|
#[derive(Debug)]
|
|
pub struct HierarchicalComposition {
|
|
levels: Vec<Vec<Box<dyn Module>>>,
|
|
level_combination_strategies: Vec<CombinationStrategy>,
|
|
config: CompositionConfig,
|
|
metadata: ModuleMetadata,
|
|
}
|
|
|
|
impl HierarchicalComposition {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
levels: Vec::new(),
|
|
level_combination_strategies: Vec::new(),
|
|
config: CompositionConfig::default(),
|
|
metadata: ModuleMetadata::default(),
|
|
}
|
|
}
|
|
|
|
pub fn add_level(mut self, modules: Vec<Box<dyn Module>>) -> Self {
|
|
self.levels.push(modules);
|
|
// Default to averaging for each level
|
|
self.level_combination_strategies
|
|
.push(CombinationStrategy::Average);
|
|
self
|
|
}
|
|
|
|
pub fn with_level_strategy(mut self, level: usize, strategy: CombinationStrategy) -> Self {
|
|
if level < self.level_combination_strategies.len() {
|
|
self.level_combination_strategies[level] = strategy;
|
|
}
|
|
self
|
|
}
|
|
|
|
pub fn with_config(mut self, config: CompositionConfig) -> Self {
|
|
self.config = config;
|
|
self
|
|
}
|
|
|
|
pub fn forward(&self, input: &Tensor) -> Result<CompositionResult> {
|
|
if self.levels.is_empty() {
|
|
return Err(ModularError::CompositionError {
|
|
message: "No levels in hierarchical composition".to_string(),
|
|
});
|
|
}
|
|
|
|
let mut current_input = input.clone();
|
|
let mut all_intermediate_outputs = Vec::new();
|
|
let mut all_active_modules = Vec::new();
|
|
let mut total_cost = 0.0;
|
|
|
|
for (level_idx, level_modules) in self.levels.iter().enumerate() {
|
|
if level_modules.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
// Process all modules in this level (parallel)
|
|
let mut level_outputs = Vec::new();
|
|
let mut level_modules_used = Vec::new();
|
|
|
|
for module in level_modules {
|
|
let output = module.forward(¤t_input)?;
|
|
level_outputs.push(output);
|
|
level_modules_used.push(module.module_id().to_string());
|
|
total_cost += module.computation_cost(current_input.dims());
|
|
}
|
|
|
|
// Combine outputs from this level
|
|
let combination_strategy = self
|
|
.level_combination_strategies
|
|
.get(level_idx)
|
|
.unwrap_or(&CombinationStrategy::Average);
|
|
|
|
let level_output = combination_strategy.combine(&level_outputs)?;
|
|
|
|
if self.config.store_intermediates {
|
|
all_intermediate_outputs.extend(level_outputs);
|
|
}
|
|
all_active_modules.extend(level_modules_used);
|
|
|
|
// Output of this level becomes input to next level
|
|
current_input = level_output;
|
|
}
|
|
|
|
let mut metadata = std::collections::HashMap::new();
|
|
metadata.insert("composition_type".to_string(), "hierarchical".to_string());
|
|
metadata.insert("num_levels".to_string(), self.levels.len().to_string());
|
|
|
|
Ok(CompositionResult {
|
|
output: current_input,
|
|
intermediate_outputs: all_intermediate_outputs,
|
|
active_modules: all_active_modules,
|
|
total_cost,
|
|
metadata,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl Module for HierarchicalComposition {
|
|
fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
|
Ok(self.forward(input)?.output)
|
|
}
|
|
|
|
fn input_dim(&self) -> usize {
|
|
self.levels
|
|
.first()
|
|
.and_then(|level| level.first())
|
|
.map(|m| m.input_dim())
|
|
.unwrap_or(0)
|
|
}
|
|
|
|
fn output_dim(&self) -> usize {
|
|
// For hierarchical composition, the output dimension depends on the last level
|
|
// and its combination strategy
|
|
if let Some(last_level) = self.levels.last() {
|
|
if let Some(last_strategy) = self.level_combination_strategies.last() {
|
|
match last_strategy {
|
|
CombinationStrategy::Concatenate => {
|
|
last_level.iter().map(|m| m.output_dim()).sum()
|
|
}
|
|
_ => last_level.first().map(|m| m.output_dim()).unwrap_or(0),
|
|
}
|
|
} else {
|
|
last_level.first().map(|m| m.output_dim()).unwrap_or(0)
|
|
}
|
|
} else {
|
|
0
|
|
}
|
|
}
|
|
|
|
fn module_type(&self) -> &str {
|
|
"hierarchical_composition"
|
|
}
|
|
|
|
fn module_id(&self) -> &str {
|
|
"hierarchical"
|
|
}
|
|
|
|
fn capabilities(&self) -> Vec<String> {
|
|
let mut all_caps = Vec::new();
|
|
for level in &self.levels {
|
|
for module in level {
|
|
all_caps.extend(module.capabilities());
|
|
}
|
|
}
|
|
all_caps.sort();
|
|
all_caps.dedup();
|
|
all_caps
|
|
}
|
|
|
|
fn metadata(&self) -> &ModuleMetadata {
|
|
&self.metadata
|
|
}
|
|
|
|
fn complexity_score(&self) -> f32 {
|
|
let mut total_complexity = 0.0;
|
|
for level in &self.levels {
|
|
let level_complexity = level
|
|
.iter()
|
|
.map(|m| m.complexity_score())
|
|
.fold(0.0, f32::max); // Max within level (parallel)
|
|
total_complexity += level_complexity; // Sum across levels (sequential)
|
|
}
|
|
total_complexity
|
|
}
|
|
|
|
fn memory_footprint(&self) -> usize {
|
|
self.levels
|
|
.iter()
|
|
.flat_map(|level| level.iter())
|
|
.map(|m| m.memory_footprint())
|
|
.sum()
|
|
}
|
|
|
|
fn is_stateful(&self) -> bool {
|
|
self.levels
|
|
.iter()
|
|
.any(|level| level.iter().any(|m| m.is_stateful()))
|
|
}
|
|
|
|
fn requires_training(&self) -> bool {
|
|
self.levels
|
|
.iter()
|
|
.any(|level| level.iter().any(|m| m.requires_training()))
|
|
}
|
|
|
|
fn parameters(&self) -> Vec<&Tensor> {
|
|
self.levels
|
|
.iter()
|
|
.flat_map(|level| level.iter())
|
|
.flat_map(|m| m.parameters())
|
|
.collect()
|
|
}
|
|
|
|
fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
|
|
self.levels
|
|
.iter_mut()
|
|
.flat_map(|level| level.iter_mut())
|
|
.flat_map(|m| m.parameters_mut())
|
|
.collect()
|
|
}
|
|
|
|
fn clone_module(&self) -> Box<dyn Module> {
|
|
let cloned_levels = self
|
|
.levels
|
|
.iter()
|
|
.map(|level| level.iter().map(|m| m.clone_module()).collect())
|
|
.collect();
|
|
|
|
Box::new(HierarchicalComposition {
|
|
levels: cloned_levels,
|
|
level_combination_strategies: self.level_combination_strategies.clone(),
|
|
config: self.config.clone(),
|
|
metadata: ModuleMetadata::default(),
|
|
})
|
|
}
|
|
|
|
fn specialize_for_task(
|
|
&self,
|
|
task_type: TaskType,
|
|
training_data: &[TestData],
|
|
) -> Result<Box<dyn Module>> {
|
|
let specialized_levels = self
|
|
.levels
|
|
.iter()
|
|
.map(|level| {
|
|
level
|
|
.iter()
|
|
.map(|m| m.specialize_for_task(task_type.clone(), training_data))
|
|
.collect::<Result<Vec<_>>>()
|
|
})
|
|
.collect::<Result<Vec<_>>>()?;
|
|
|
|
Ok(Box::new(HierarchicalComposition {
|
|
levels: specialized_levels,
|
|
level_combination_strategies: self.level_combination_strategies.clone(),
|
|
config: self.config.clone(),
|
|
metadata: ModuleMetadata::default(),
|
|
}))
|
|
}
|
|
|
|
fn continually_specialize(&mut self, training_data: &[TestData], task_id: usize) -> Result<()> {
|
|
for level in &mut self.levels {
|
|
for module in level {
|
|
module.continually_specialize(training_data, task_id)?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|