Optimization API Reference¶

Overview¶

The Opifex optimization module provides full optimization algorithms and meta-learning approaches for scientific computing, including production optimization, learn-to-optimize algorithms, control systems, and quantum-aware optimization.

Core Optimization Components¶

Meta-Optimization Framework¶

Advanced meta-optimization algorithms that learn to optimize across families of related problems.

Meta-optimization algorithms for scientific machine learning.

This package implements meta-learning approaches to optimization including learn-to-optimize (L2O) algorithms, adaptive learning rate scheduling, warm-starting strategies, and performance monitoring. All implementations follow FLAX NNX patterns and are designed for scientific computing applications.

Key Features

Learn-to-optimize (L2O) meta-learning algorithms
Adaptive learning rate scheduling with multiple strategies
Warm-starting based on problem similarity
Performance monitoring and analytics
Quantum-aware optimization adaptations
Integration with existing training infrastructure

Author: Opifex Framework Team Date: December 2024 License: MIT

MetaOptimizerConfig `dataclass` ¶

MetaOptimizerConfig(*, meta_algorithm: str = 'l2o', base_optimizer: str = 'adam', meta_learning_rate: float = 0.0001, adaptation_steps: int = 10, warm_start_strategy: str = 'previous_params', performance_tracking: bool = True, memory_efficient: bool = True, quantum_aware: bool = False, scf_adaptation: bool = False, energy_convergence_tracking: bool = False, chemical_accuracy_target: float = 0.001)

Configuration for meta-optimization algorithms.

This configuration class defines all parameters for meta-optimization including algorithm selection, adaptation strategies, and performance monitoring settings.

Attributes:

Name	Type	Description
`meta_algorithm`	`str`	Meta-optimization algorithm ('l2o', 'adaptive_lr', 'warm_start')
`base_optimizer`	`str`	Base optimizer to enhance ('adam', 'sgd', 'rmsprop', 'adamw')
`meta_learning_rate`	`float`	Learning rate for meta-parameters
`adaptation_steps`	`int`	Number of steps for adaptation
`warm_start_strategy`	`str`	Strategy for warm-starting ('previous_params', 'similar_problems')
`performance_tracking`	`bool`	Enable performance monitoring
`memory_efficient`	`bool`	Use memory-efficient implementations
`quantum_aware`	`bool`	Enable quantum-specific adaptations
`scf_adaptation`	`bool`	Enable SCF convergence acceleration
`energy_convergence_tracking`	`bool`	Track energy convergence for quantum systems
`chemical_accuracy_target`	`float`	Target chemical accuracy (kcal/mol)

MetaOptimizer ¶

MetaOptimizer(config: MetaOptimizerConfig, *, rngs: Rngs)

Integrated meta-optimization system.

This class provides a complete meta-optimization system that integrates learn-to-optimize algorithms, adaptive learning rate scheduling, warm-starting strategies, and performance monitoring.

Attributes:

Name	Type	Description
`config`		Meta-optimizer configuration
`l2o_engine`		Learn-to-optimize engine
`learning_rate_scheduler`		Adaptive learning rate scheduler
`warm_start_strategy`		Warm-starting strategy
`performance_monitor`		Performance monitoring system
`current_step`		Current optimization step

Parameters:

Name	Type	Description	Default
`config`	`MetaOptimizerConfig`	Meta-optimizer configuration	required
`rngs`	`Rngs`	Random number generators	required

init_optimizer_state ¶

init_optimizer_state(params: Array) -> Any

Initialize optimizer state.

Parameters:

Name	Type	Description	Default
`params`	`Array`	Initial parameters	required

Returns:

Type	Description
`Any`	Initial optimizer state

step ¶

step(loss_fn: Callable[[Array], Array], params: Array, opt_state: Any, step: int) -> tuple[Array, Any, dict[str, Any]]

Perform single meta-optimization step.

Parameters:

Name	Type	Description	Default
`loss_fn`	`Callable[[Array], Array]`	Loss function to optimize	required
`params`	`Array`	Current parameters	required
`opt_state`	`Any`	Current optimizer state	required
`step`	`int`	Current step number	required

Returns:

Type	Description
`tuple[Array, Any, dict[str, Any]]`	Tuple of (new_params, new_opt_state, meta_info)

store_optimization_result ¶

store_optimization_result(params: Array, problem_features: Array) -> None

Store optimization result for future warm-starting.

Parameters:

Name	Type	Description	Default
`params`	`Array`	Final optimized parameters	required
`problem_features`	`Array`	Features characterizing the problem	required

get_warm_start_params ¶

get_warm_start_params(current_problem_features: Array, target_shape: tuple[int, ...]) -> Array

Get warm-start parameters for new problem.

Parameters:

Name	Type	Description	Default
`current_problem_features`	`Array`	Features of current problem	required
`target_shape`	`tuple[int, ...]`	Target shape for parameters	required

Returns:

Type	Description
`Array`	Warm-start parameters

quantum_step ¶

quantum_step(energy_fn: Callable[[Array], Array], orbital_coeffs: Array, opt_state: Any, scf_context: dict[str, Any], step: int) -> tuple[Array, Any, dict[str, Any]]

Perform quantum-aware meta-optimization step.

Parameters:

Name	Type	Description	Default
`energy_fn`	`Callable[[Array], Array]`	Energy function to minimize	required
`orbital_coeffs`	`Array`	Current orbital coefficients	required
`opt_state`	`Any`	Current optimizer state	required
`scf_context`	`dict[str, Any]`	SCF iteration context	required
`step`	`int`	Current step number	required

Returns:

Type	Description
`tuple[Array, Any, dict[str, Any]]`	Tuple of (new_coeffs, new_opt_state, quantum_info)

PerformanceMonitor ¶

PerformanceMonitor(metrics: list[str] | None = None, window_size: int = 100, tracking_frequency: int = 1, convergence_tolerance: float = 1e-06, convergence_patience: int = 10, analytics_enabled: bool = False, quantum_aware: bool = False)

Performance monitoring and analytics for meta-optimization.

This class provides full performance monitoring capabilities including metric tracking, convergence detection, and performance analytics for optimization algorithms.

Attributes:

Name	Type	Description
`metrics`		List of metrics to track
`window_size`		Size of rolling window for metrics
`tracking_frequency`		Frequency of metric updates
`convergence_tolerance`		Tolerance for convergence detection
`convergence_patience`		Patience for convergence detection
`analytics_enabled`		Enable detailed analytics
`quantum_aware`		Enable quantum-specific metrics

Parameters:

Name	Type	Description	Default
`metrics`	`list[str] \| None`	List of metrics to track	`None`
`window_size`	`int`	Rolling window size	`100`
`tracking_frequency`	`int`	How often to update metrics	`1`
`convergence_tolerance`	`float`	Tolerance for convergence	`1e-06`
`convergence_patience`	`int`	Patience for convergence detection	`10`
`analytics_enabled`	`bool`	Enable detailed analytics	`False`
`quantum_aware`	`bool`	Enable quantum metrics	`False`

update_metrics ¶

update_metrics(step: int, **metric_values: float) -> None

Update tracked metrics.

Parameters:

Name	Type	Description	Default
`step`	`int`	Current optimization step	required
`**metric_values`	`float`	Metric values to update	`{}`

get_metric_history ¶

get_metric_history(metric: str) -> list[float]

Get history of a specific metric.

Parameters:

Name	Type	Description	Default
`metric`	`str`	Metric name	required

Returns:

Type	Description
`list[float]`	List of metric values

check_convergence ¶

check_convergence(metric: str) -> bool

Check if a metric has converged.

Parameters:

Name	Type	Description	Default
`metric`	`str`	Metric name	required

Returns:

Type	Description
`bool`	True if metric has converged

get_performance_analytics ¶

get_performance_analytics() -> dict[str, Any]

Get full performance analytics.

Returns:

Type	Description
`dict[str, Any]`	Dictionary containing performance analytics

get_quantum_analytics ¶

get_quantum_analytics() -> dict[str, Any]

Get quantum-specific performance analytics.

Returns:

Type	Description
`dict[str, Any]`	Dictionary containing quantum analytics

LearnToOptimize ¶

LearnToOptimize(meta_network_layers: list[int] | None = None, base_optimizer: str = 'adam', meta_learning_rate: float = 0.0001, unroll_steps: int = 20, adaptive_step_size: bool = False, quantum_aware: bool = False, scf_integration: bool = False, *, rngs: Rngs)

Bases: Module

Learn-to-optimize (L2O) meta-learning system.

This class implements learn-to-optimize algorithms that use neural networks to learn optimization strategies from data. The meta-network learns to predict good parameter updates based on gradient information and optimization history.

Attributes:

Name	Type	Description
`meta_network`		Neural network for learning optimization rules
`base_optimizer`		Base optimization algorithm
`meta_learning_rate`		Learning rate for meta-network training
`unroll_steps`		Number of unrolling steps for meta-gradient computation
`adaptive_step_size`		Enable adaptive step size learning
`quantum_aware`		Enable quantum-specific adaptations
`scf_integration`		Enable SCF convergence acceleration

Parameters:

Name	Type	Description	Default
`meta_network_layers`	`list[int] \| None`	Architecture of meta-network	`None`
`base_optimizer`	`str`	Base optimizer to enhance	`'adam'`
`meta_learning_rate`	`float`	Learning rate for meta-network training	`0.0001`
`unroll_steps`	`int`	Number of unroll steps for meta-gradients	`20`
`adaptive_step_size`	`bool`	Enable adaptive step size learning	`False`
`quantum_aware`	`bool`	Enable quantum-specific optimizations	`False`
`scf_integration`	`bool`	Enable SCF convergence acceleration	`False`
`rngs`	`Rngs`	Random number generators for initialization	required

compute_update ¶

compute_update(gradient: Array, previous_updates: Array, loss_history: Array | None = None) -> Array

Compute parameter update using meta-network.

Parameters:

Name	Type	Description	Default
`gradient`	`Array`	Current gradient	required
`previous_updates`	`Array`	History of previous updates	required
`loss_history`	`Array \| None`	History of loss values	`None`

Returns:

Type	Description
`Array`	Predicted parameter update

compute_meta_gradients ¶

compute_meta_gradients(loss_fn: Callable[[Array], Array], initial_params: Array) -> dict[str, Array]

Compute meta-gradients for meta-network training.

Parameters:

Name	Type	Description	Default
`loss_fn`	`Callable[[Array], Array]`	Loss function for optimization problem	required
`initial_params`	`Array`	Initial parameters for optimization	required

Returns:

Type	Description
`dict[str, Array]`	Meta-gradients for meta-network parameters

compute_adaptive_update ¶

compute_adaptive_update(gradient: Array, previous_updates: Array) -> Array

Compute adaptive parameter update.

compute_quantum_update ¶

compute_quantum_update(orbital_params: Array, scf_history: Array) -> Array

Compute quantum-aware parameter update.

Parameters:

Name	Type	Description	Default
`orbital_params`	`Array`	Orbital coefficient parameters	required
`scf_history`	`Array`	SCF convergence history	required

Returns:

Type	Description
`Array`	Quantum-adapted parameter update

AdaptiveLearningRateScheduler ¶

AdaptiveLearningRateScheduler(schedule_type: str = 'cosine_annealing', initial_lr: float = 0.001, final_lr: float = 1e-06, adaptation_period: int = 100, warmup_steps: int = 0, patience: int = 5, factor: float = 0.5, min_lr: float = 1e-08, **kwargs: Any)

Adaptive learning rate scheduling for meta-optimization.

This class implements various adaptive learning rate scheduling strategies including cosine annealing, performance-based adaptation, and quantum-aware scheduling for scientific applications.

Attributes:

Name	Type	Description
`schedule_type`		Type of scheduling algorithm
`initial_lr`		Initial learning rate
`final_lr`		Final learning rate (for annealing schedules)
`adaptation_period`		Period for adaptation cycles
`warmup_steps`		Number of warmup steps
`patience`		Patience for performance-based adaptation
`factor`		Factor for learning rate reduction
`min_lr`		Minimum learning rate

Parameters:

Name	Type	Description	Default
`schedule_type`	`str`	Type of scheduling ('cosine_annealing', 'performance_based', 'quantum_aware')	`'cosine_annealing'`
`initial_lr`	`float`	Initial learning rate	`0.001`
`final_lr`	`float`	Final learning rate	`1e-06`
`adaptation_period`	`int`	Period for complete adaptation cycle	`100`
`warmup_steps`	`int`	Number of warmup steps	`0`
`patience`	`int`	Patience for performance-based adaptation	`5`
`factor`	`float`	Reduction factor for learning rate	`0.5`
`min_lr`	`float`	Minimum allowed learning rate	`1e-08`
`**kwargs`	`Any`	Additional scheduler-specific parameters	`{}`

get_learning_rate ¶

get_learning_rate(step: int) -> Array

Get learning rate for current step.

Parameters:

Name	Type	Description	Default
`step`	`int`	Current optimization step	required

Returns:

Type	Description
`Array`	Learning rate for current step

adapt_from_performance ¶

adapt_from_performance(loss_history: list[float]) -> float

Adapt learning rate based on performance history.

Parameters:

Name	Type	Description	Default
`loss_history`	`list[float]`	Recent loss values	required

Returns:

Type	Description
`float`	Adapted learning rate

adapt_from_quantum_metrics ¶

adapt_from_quantum_metrics(scf_errors: list[float], energy_changes: list[float]) -> float

Adapt learning rate based on quantum mechanical metrics.

Parameters:

Name	Type	Description	Default
`scf_errors`	`list[float]`	SCF convergence errors	required
`energy_changes`	`list[float]`	Energy change magnitudes	required

Returns:

Type	Description
`float`	Quantum-adapted learning rate

WarmStartingStrategy ¶

WarmStartingStrategy(strategy_type: str = 'parameter_transfer', similarity_threshold: float = 0.8, adaptation_steps: int = 5, memory_size: int = 10, adaptation_ratio: float = 0.9, similarity_metric: str = 'cosine', min_similarity: float = 0.7)

Warm-starting strategies for optimization acceleration.

This class implements various warm-starting strategies to accelerate optimization by leveraging information from previous optimizations or similar problems.

Attributes:

Name	Type	Description
`strategy_type`		Type of warm-starting strategy
`similarity_threshold`		Threshold for problem similarity
`adaptation_steps`		Steps for parameter adaptation
`memory_size`		Size of optimization memory
`adaptation_ratio`		Ratio for optimizer state adaptation

Parameters:

Name	Type	Description	Default
`strategy_type`	`str`	Strategy type ('parameter_transfer', 'optimizer_state_transfer', 'molecular_similarity')	`'parameter_transfer'`
`similarity_threshold`	`float`	Threshold for considering problems similar	`0.8`
`adaptation_steps`	`int`	Number of adaptation steps	`5`
`memory_size`	`int`	Maximum number of previous optimizations to remember	`10`
`adaptation_ratio`	`float`	Ratio for adapting previous states	`0.9`
`similarity_metric`	`str`	Metric for similarity computation	`'cosine'`
`min_similarity`	`float`	Minimum similarity for warm-starting	`0.7`

get_warm_start_params ¶

get_warm_start_params(previous_params: Array, current_problem_features: Array) -> Array

Get warm-start parameters based on parameter transfer.

Parameters:

Name	Type	Description	Default
`previous_params`	`Array`	Parameters from previous optimization	required
`current_problem_features`	`Array`	Features of current problem	required

Returns:

Type	Description
`Array`	Warm-start parameters for current problem

adapt_optimizer_state ¶

adapt_optimizer_state(previous_opt_state: dict[str, Any]) -> dict[str, Any]

Adapt optimizer state for warm-starting.

Parameters:

Name	Type	Description	Default
`previous_opt_state`	`dict[str, Any]`	Previous optimizer state	required

Returns:

Type	Description
`dict[str, Any]`	Adapted optimizer state

get_molecular_warm_start ¶

get_molecular_warm_start(previous_fingerprints: Array, previous_params: Array, current_fingerprint: Array) -> Array

Get warm-start parameters based on molecular similarity.

Parameters:

Name	Type	Description	Default
`previous_fingerprints`	`Array`	Fingerprints of previous molecules	required
`previous_params`	`Array`	Parameters for previous molecules	required
`current_fingerprint`	`Array`	Fingerprint of current molecule	required

Returns:

Type	Description
`Array`	Warm-start parameters based on most similar molecule

Production Optimization¶

Optimization systems for deployment and scaling in production environments.

Production optimisation for the Opifex framework.

The :class:HybridPerformancePlatform combines adaptive JIT compilation (:class:AdaptiveJAXOptimizer), GPU memory-pool planning (:class:IntelligentGPUMemoryManager), and physics/numerical validation (:class:~opifex.optimization.scientific_integration.ScientificComputingIntegrator) into a single production-optimisation pass. Serving telemetry, autoscaling, and edge/deployment orchestration are out of scope (owned by external infrastructure such as KServe / Ray Serve / k8s HPA / Prometheus).

CallableModule ¶

Bases: Protocol

Protocol for callable modules.

OptimizationStrategy ¶

Bases: Enum

JIT optimization strategies for different workload patterns.

WorkloadProfile `dataclass` ¶

WorkloadProfile(batch_size: int, sequence_length: int, memory_footprint: float, compute_intensity: float, latency_requirement: float, throughput_requirement: float, model_complexity: str)

Profiling data for production workloads.

PerformanceMetrics `dataclass` ¶

PerformanceMetrics(*, latency_ms: float, throughput_rps: float, memory_usage_gb: float, improvement_factor: float)

Measured performance metrics for an optimized model.

All fields are directly measured. GPU utilization and energy efficiency are intentionally omitted: measuring them requires device/power telemetry (e.g. NVML) that is not a dependency of this framework, so they cannot be reported here as measurements.

OptimizedModel `dataclass` ¶

OptimizedModel(*, model: Module, optimization_type: OptimizationStrategy, performance_metrics: PerformanceMetrics, optimization_metadata: dict[str, Any])

Container for optimized model with performance metadata.

AdaptiveJAXOptimizer ¶

AdaptiveJAXOptimizer(performance_threshold: float = 1.1, memory_efficiency_target: float = 0.85, cache_size: int = 100)

Bases: Module

Adaptive JIT optimization for JAX-based neural operators.

This class implements intelligent JIT compilation strategies based on workload patterns, providing optimal performance for production deployments.

Parameters:

Name	Type	Description	Default
`performance_threshold`	`float`	Minimum performance improvement factor to accept optimization	`1.1`
`memory_efficiency_target`	`float`	Target memory efficiency (0-1 scale)	`0.85`
`cache_size`	`int`	Number of optimization strategies to cache	`100`

analyze_workload_patterns ¶

analyze_workload_patterns(workload: WorkloadProfile) -> OptimizationStrategy

Analyze workload to select optimal optimization strategy.

apply_aggressive_kernel_fusion ¶

apply_aggressive_kernel_fusion(model: Module) -> Module

Apply aggressive kernel fusion for compute-intensive workloads.

apply_memory_optimization ¶

apply_memory_optimization(model: Module) -> Module

Apply memory optimization for large models.

apply_latency_optimization ¶

apply_latency_optimization(model: Module) -> Module

Apply latency optimization for real-time inference.

apply_balanced_optimization ¶

apply_balanced_optimization(model: Module) -> Module

Apply balanced optimization for general workloads.

benchmark_model_performance ¶

benchmark_model_performance(model: Module, workload: WorkloadProfile, improvement_factor: float = 1.0) -> PerformanceMetrics

Benchmark a single model on the given workload.

Latency is measured by timing repeated forward passes; throughput is derived from that latency. Memory usage is estimated from the workload footprint. improvement_factor defaults to 1.0 (no baseline to compare against in a single-model benchmark) and should be supplied by the caller when a measured baseline latency is available.

Parameters:

Name	Type	Description	Default
`model`	`Module`	Model to benchmark.	required
`workload`	`WorkloadProfile`	Workload profile defining batch size and footprint.	required
`improvement_factor`	`float`	Measured speedup relative to a baseline model (baseline_latency / this_latency); `1.0` when no baseline.	`1.0`

Returns:

Type	Description
`PerformanceMetrics`	Measured performance metrics for the model.

optimize_neural_operator ¶

optimize_neural_operator(model: Module, workload: WorkloadProfile) -> OptimizedModel

Optimize neural operator for production workload.

Parameters:

Name	Type	Description	Default
`model`	`Module`	Neural operator model to optimize	required
`workload`	`WorkloadProfile`	Workload profile for optimization	required

Returns:

Type	Description
`OptimizedModel`	OptimizedModel with performance improvements

IntelligentGPUMemoryManager ¶

IntelligentGPUMemoryManager(fragmentation_threshold: float = 0.15, gc_trigger_threshold: float = 0.85, pool_sizes: dict[str, tuple[int, int]] | None = None)

Bases: Module

Advanced GPU memory management for production workloads.

Implements intelligent allocation, fragmentation prevention, and multi-model inference optimization.

Parameters:

Name	Type	Description	Default
`fragmentation_threshold`	`float`	Maximum acceptable fragmentation (0-1)	`0.15`
`gc_trigger_threshold`	`float`	Memory usage threshold to trigger GC (0-1)	`0.85`
`pool_sizes`	`dict[str, tuple[int, int]] \| None`	Memory pool sizes as {pool_name: (min_size_mb, max_size_mb)}	`None`

select_memory_pool ¶

select_memory_pool(size_mb: float) -> str

Select appropriate memory pool for allocation size.

estimate_model_memory_usage ¶

estimate_model_memory_usage(model: Module, batch_size: int) -> float

Estimate memory usage for model inference (in MB).

optimize_multi_model_allocation ¶

optimize_multi_model_allocation(models: list[tuple[Module, int]]) -> dict[str, Any]

Optimize memory allocation for multiple concurrent models.

Parameters:

Name	Type	Description	Default
`models`	`list[tuple[Module, int]]`	List of (model, batch_size) tuples	required

Returns:

Type	Description
`dict[str, Any]`	Allocation plan with memory optimization strategy

HybridPerformancePlatform ¶

HybridPerformancePlatform(jit_optimizer: AdaptiveJAXOptimizer | None = None, memory_manager: IntelligentGPUMemoryManager | None = None, scientific_integrator: ScientificComputingIntegrator | None = None, physics_domain: PhysicsDomain = GENERAL, target_latency_ms: float = 0.5)

Bases: Module

Production optimisation orchestrator: JIT, GPU memory, and scientific validation.

Combines the three genuine optimisation components — :class:AdaptiveJAXOptimizer (jax.jit kernel fusion), :class:IntelligentGPUMemoryManager (memory-pool planning), and :class:~opifex.optimization.scientific_integration.ScientificComputingIntegrator (physics / numerical validation) — into a single production-optimisation pass. It does not perform serving telemetry or autoscaling; those concerns belong to external infrastructure (KServe / Ray Serve / k8s HPA / Prometheus).

optimize_for_production ¶

optimize_for_production(model: Module, workload: WorkloadProfile) -> OptimizedModel

Full production optimisation for a model: JIT, memory, and scientific validation.

get_model_input_features ¶

get_model_input_features(model: Module) -> int

Extract the correct input feature dimension from a model.

This utility function provides a robust way to determine the expected input dimension for any model, preventing dimension mismatch errors.

Parameters:

Name	Type	Description	Default
`model`	`Module`	The model to inspect	required

Returns:

Name	Type	Description
`int`	`int`	The input feature dimension

Raises:

Type	Description
`ValueError`	If input features cannot be determined

Scientific Computing Integration¶

Physics-aware optimization with scientific validation and benchmarking.

Scientific computing integration for Opifex production optimization.

This module implements physics-informed optimization, numerical validation, and conservation checking for the Version 7.4 Production Optimization system.

Part of: Hybrid Performance Platform + Intelligent Edge + Adaptive Optimization

PhysicsDomain ¶

Bases: Enum

Scientific computing domains.

PhysicsMetrics `dataclass` ¶

PhysicsMetrics(*, domain: PhysicsDomain, conservation_violations: dict[ConservationLaw, float] = dict(), symmetry_preservation: float = 0.0, numerical_stability: float = 0.0, energy_conservation_error: float = 0.0, momentum_conservation_error: float = 0.0, mass_conservation_error: float = 0.0, unitarity_preservation: float = 0.0, thermodynamic_consistency: float = 0.0, boundary_condition_accuracy: float = 0.0)

Physics-specific performance metrics.

NumericalValidationResult `dataclass` ¶

NumericalValidationResult(*, is_valid: bool, precision_score: float, stability_score: float, convergence_rate: float, condition_number: float, validation_errors: list[str] = list(), recommendations: list[str] = list())

Result of numerical validation.

ConservationCheckResult `dataclass` ¶

ConservationCheckResult(*, law: ConservationLaw, is_conserved: bool, violation_magnitude: float, tolerance: float, relative_error: float, time_evolution_consistency: bool = True)

Result of conservation law checking.

ScientificBenchmarkResult `dataclass` ¶

ScientificBenchmarkResult(*, benchmark_name: str, domain: PhysicsDomain, accuracy_score: float, reference_value: float, computed_value: float, relative_error: float, meets_accuracy_threshold: bool, chemical_accuracy: bool = False)

Result of scientific benchmark validation.

PhysicsProfilerProtocol ¶

Bases: Protocol

Protocol for physics profiling implementations.

profile_physics_metrics ¶

profile_physics_metrics(model_output: ndarray, reference_data: dict[str, Any]) -> PhysicsMetrics

Profile physics-specific metrics.

validate_domain_constraints ¶

validate_domain_constraints(model_output: ndarray, domain: PhysicsDomain) -> bool

Validate domain-specific constraints.

PhysicsProfiler ¶

PhysicsProfiler(domain: PhysicsDomain, validation_tolerances: dict[str, float] | None = None)

Physics-informed profiler for domain-specific optimization.

profile_physics_metrics ¶

profile_physics_metrics(model_output: ndarray, reference_data: dict[str, Any], time_series: list[ndarray] | None = None) -> PhysicsMetrics

Profile full physics-specific metrics.

NumericalValidator ¶

NumericalValidator(precision_threshold: float = 1e-06, stability_threshold: float = 0.001)

Numerical precision and stability validator.

validate_numerical_precision ¶

validate_numerical_precision(computed_values: ndarray, reference_values: ndarray) -> NumericalValidationResult

Validate numerical precision against reference values.

check_conservation_law ¶

check_conservation_law(computed_quantity: ndarray, reference_quantity: ndarray, law: ConservationLaw, tolerance: float | None = None) -> ConservationCheckResult

Check specific conservation law.

ScientificBenchmarkValidator ¶

ScientificBenchmarkValidator(domain: PhysicsDomain)

Validator for scientific computing benchmarks.

validate_benchmark ¶

validate_benchmark(benchmark_name: str, computed_value: float, reference_value: float, accuracy_type: str = 'relative_error') -> ScientificBenchmarkResult

Validate against a specific benchmark.

validate_multiple_benchmarks ¶

validate_multiple_benchmarks(benchmarks: dict[str, tuple[float, float]]) -> list[ScientificBenchmarkResult]

Validate multiple benchmarks.

ScientificComputingIntegrator ¶

ScientificComputingIntegrator(domain: PhysicsDomain, physics_profiler: PhysicsProfiler | None = None, numerical_validator: NumericalValidator | None = None, benchmark_validator: ScientificBenchmarkValidator | None = None)

Main integrator for scientific computing optimization.

comprehensive_scientific_validation ¶

comprehensive_scientific_validation(model_output: ndarray, reference_data: dict[str, Any], benchmarks: dict[str, tuple[float, float]] | None = None, time_series: list[ndarray] | None = None) -> dict[str, Any]

Perform full scientific validation.

optimize_for_scientific_accuracy ¶

optimize_for_scientific_accuracy(model_output: ndarray, validation_results: dict[str, Any]) -> dict[str, Any]

Generate optimization recommendations based on scientific validation.

translation_invariance_error ¶

translation_invariance_error(field: Array, shifts: tuple[int, ...], axes: tuple[int, ...]) -> Array

Relative L2 discrepancy between a field and its periodic translation.

A translation-invariant field is unchanged by a periodic spatial shift, so shifting it and measuring the relative L2 difference yields ~0 for a constant or periodic field and a positive value for a field with spatial structure.

Parameters:

Name	Type	Description	Default
`field`	`Array`	Spatial field to test for translation invariance.	required
`shifts`	`tuple[int, ...]`	Integer roll amount applied along each axis in `axes`.	required
`axes`	`tuple[int, ...]`	Axes along which the periodic shift is applied.	required

Returns:

Type	Description
`Array`	Relative L2 discrepancy between the field and its shifted copy.

Learn-to-Optimize (L2O) Algorithms¶

Advanced neural optimization methods that achieve significant speedups on learned problem families.

Tasks and the optimiser interface¶

Objective-carrying Task/TaskFamily abstractions and the shared stateful Optimizer interface (with the optax-wrapped baseline family).

Core learn-to-optimize abstractions: tasks and the optimizer interface.

A learned optimizer is meta-trained to minimise a distribution of objectives. The objective is carried by a :class:Task (an init for the optimisee parameters plus a loss), and a :class:TaskFamily samples tasks for meta-generalisation. This mirrors the canonical design in Google's learned_optimization library (learned_optimization/tasks/base.py); see Andrychowicz et al. 2016 (arXiv:1606.04474) for the original learning-to-learn formulation.

The key contrast with the previous opifex L2O code is that the objective lives on the task: optimisers and meta-trainers close over task.loss rather than guessing a placeholder. Task.normalizer maps a raw loss onto a comparable scale so a meta-loss aggregated across differently-conditioned tasks is not dominated by the worst-scaled task (learned_optimization/tasks/base.py normalizer).

Task ¶

Bases: ABC

A single optimisation objective: init the params, evaluate loss.

Subclasses implement :meth:init (sample the initial optimisee parameters) and :meth:loss (a scalar objective). :meth:normalizer defaults to the identity and should be overridden when meta-training mixes tasks of very different loss scales.

init `abstractmethod` ¶

init(key: Array) -> PyTree

Sample the initial optimisee parameter pytree.

loss `abstractmethod` ¶

loss(params: PyTree, key: Array) -> Float[Array, '']

Return the scalar objective at params (key for stochastic tasks).

normalizer ¶

normalizer(loss: Float[Array, '']) -> Float[Array, '']

Map a raw loss onto a comparable scale (identity by default).

loss_and_grad ¶

loss_and_grad(params: PyTree, key: Array) -> tuple[Float[Array, ''], PyTree]

Convenience: value-and-gradient of :meth:loss w.r.t. params.

TaskFamily ¶

Bases: ABC

A distribution over :class:Task instances for meta-generalisation.

sample `abstractmethod` ¶

sample(key: Array) -> Task

Draw a task from the family.

single_task_to_family ¶

single_task_to_family(task: Task) -> TaskFamily

Lift a fixed :class:Task into a :class:TaskFamily (mirrors the reference).

Meta-training always consumes a family; this adapts a single task for the case where no task distribution is needed (e.g. overfitting a learned optimiser to one problem).

The optimiser interface shared by hand-designed and learned optimisers.

Mirrors learned_optimization/optimizers/base.py: an :class:Optimizer threads an opaque state through init / update / get_params. Both an optax-wrapped hand-designed optimiser (:class:OptaxOptimizer) and a meta-learned optimiser fit this interface, so they are interchangeable at every call site. init accepts num_steps so horizon-aware optimisers (e.g. learned optimisers conditioning on training fraction) know the unroll length; update accepts the current loss because learned optimisers may consume it.

OptaxOptState ¶

Bases: PyTreeNode

Pytree state for :class:OptaxOptimizer (params + optax state + step count).

Optimizer ¶

Bases: ABC, Generic[StateT]

Stateful optimiser interface (object-aware, not an optax.GradientTransformation).

The state is opaque to callers; use :meth:get_params to read the current optimisee parameters. num_steps (at :meth:init) is the planned unroll length, surfaced for horizon-aware optimisers; loss (at :meth:update) is the current objective value, consumed by learned optimisers and ignored by hand-designed ones.

init `abstractmethod` ¶

init(params: PyTree, *, num_steps: int | None = None, key: Array | None = None) -> StateT

Initialise optimiser state for the given starting params.

update `abstractmethod` ¶

update(state: StateT, grad: PyTree, *, loss: Array | None = None) -> StateT

Apply one optimiser step, returning the new state.

get_params `abstractmethod` ¶

get_params(state: StateT) -> PyTree

Read the current optimisee parameters from state.

OptaxOptimizer ¶

OptaxOptimizer(transformation: GradientTransformation)

Bases: Optimizer[OptaxOptState]

Wraps any optax.GradientTransformation as an :class:Optimizer.

This is the hand-designed baseline family (SGD/Adam/...). num_steps/key are accepted for interface uniformity and ignored; loss is ignored (optax updates depend only on the gradient and optimiser state).

init ¶

init(params: PyTree, *, num_steps: int | None = None, key: Array | None = None) -> OptaxOptState

Initialise optax state around params.

update ¶

update(state: OptaxOptState, grad: PyTree, *, loss: Array | None = None) -> OptaxOptState

Apply the optax update and advance the iteration counter.

get_params ¶

get_params(state: OptaxOptState) -> PyTree

Return the current parameters.

Concrete tasks¶

QuadraticTaskFamily (convex smoke task) and MLPTaskFamily (the non-convex small-MLP training showcase task).

Concrete optimisation tasks and task families for L2O meta-training/meta-test.

QuadraticTask is the canonical L2O smoke task (a strictly convex quadratic with a known optimum; cf. learned_optimization/tasks/quadratics.py). QuadraticTaskFamily samples quadratics with varied conditioning so a meta-trained optimiser must generalise across loss landscapes of different curvature/scale — the diversity that makes a meta-test on held-out tasks meaningful (Wichrowska et al. 2017, arXiv:1703.04813).

MLPTask is the canonical L2O showcase task: a small multilayer perceptron trained by the inner optimiser. It mirrors the MLPTask used throughout Google's learned_optimization tutorials (docs/notebooks/no_dependency_learned_optimizer) — a genuinely non-convex neural-network training objective, which is the regime where learned optimisers demonstrably beat fixed-hyperparameter baselines (Metz et al. 2020, arXiv:2009.11243). To stay self-contained (no dataset dependency) the data is a synthetic teacher-student regression: a random teacher MLP generates targets from Gaussian inputs and the student fits them with MSE, so the global optimum is realisable yet the landscape is non-convex.

QuadraticTask `dataclass` ¶

QuadraticTask(matrix: Array, optimum: Array, init_scale: float = 1.0, loss_scale: Array | float = 1.0)

Bases: Task

Strictly convex quadratic f(x) = 0.5 (x - x*)^T A (x - x*) with SPD A.

Shapes: matrix is (dim, dim) SPD, optimum (x*) is (dim,). The unique minimiser is optimum with f(x*) = 0. The objective is deterministic, so the per-call key is unused.

init ¶

init(key: Array) -> Array

Sample a (dim,) starting point offset from the optimum by init_scale.

loss ¶

loss(params: Array, key: Array) -> Array

Evaluate the scalar quadratic objective at params (dim,).

normalizer ¶

normalizer(loss: Array) -> Array

Scale by the expected initial loss so normalised loss is O(1) across tasks.

QuadraticTaskFamily `dataclass` ¶

QuadraticTaskFamily(dim: int, init_scale: float = 1.0, max_log_condition: float = 3.0)

Bases: TaskFamily

Quadratics with random optima and random SPD curvature of varied conditioning.

Each draw builds A = Q diag(eigs) Q^T with a random orthogonal Q and eigenvalues log-uniformly spread over a sampled condition number in [1, 10**max_log_condition] — so tasks differ in both scale and conditioning.

sample ¶

sample(key: Array) -> QuadraticTask

Draw a quadratic task: random SPD A, random optimum, matched loss scale.

MLPTask `dataclass` ¶

MLPTask(inputs: Array, targets: Array, teacher_params: MLPParams, layer_sizes: tuple[int, ...], init_scale: float, loss_scale: Array, batch_size: int)

Bases: Task

Minibatch MLP regression onto a per-task teacher-generated dataset.

The student MLP (architecture layer_sizes) is fit by MSE to targets produced by a random teacher MLP on inputs. Each :meth:loss call draws a fresh batch_size minibatch using the supplied key (mirroring the reference MLPTask, which consumes a new minibatch per step), so the gradients are stochastic — the regime where a learned optimiser's implicit learning-rate schedule beats a fixed-step baseline. The objective is non-convex in the student weights; the optimum (loss 0) is realised at teacher_params.

init ¶

init(key: Array) -> MLPParams

Sample initial student parameters with the small-init convention.

loss ¶

loss(params: MLPParams, key: Array) -> Array

Mean-squared error on a fresh batch_size minibatch drawn with key.

normalizer ¶

normalizer(loss: Array) -> Array

Scale by the expected initial loss so normalised loss is O(1) across tasks.

MLPTaskFamily `dataclass` ¶

MLPTaskFamily(input_dim: int, hidden_dim: int, output_dim: int, num_data: int = 512, batch_size: int = 32, init_scale: float = 0.02, teacher_scale: float = 1.0)

Bases: TaskFamily

Teacher-student MLP regression tasks: random teacher + synthetic Gaussian data per draw.

Each draw samples a fresh teacher MLP and a fresh dataset, so a meta-trained optimiser must generalise across many non-convex training problems rather than memorise one (the diversity that makes the held-out meta-test meaningful). Gradients are stochastic minibatch gradients.

layer_sizes `property` ¶

layer_sizes: tuple[int, ...]

The student/teacher architecture (input, hidden, output).

sample ¶

sample(key: Array) -> MLPTask

Draw a teacher-student regression task with fresh teacher and inputs.

Learned optimisers¶

Coordinatewise learned optimisers (the per-parameter MLP of Metz et al. 2020) and their input features.

Per-parameter input features for learned optimisers.

Faithful re-implementation of the feature primitives in Google's learned_optimization (learned_optimizers/common.py and learned_optimizers/mlp_lopt.py). A coordinatewise learned optimiser consumes, per scalar parameter: multi-timescale momentum and RMS EMAs of the gradient, the gradient and parameter themselves (second-moment-normalised across the tensor), and a tanh embedding of the iteration (training-fraction awareness). Feature design follows Metz et al. 2020 (arXiv:2009.11243).

These are pure functions on arrays; learned.py composes them into the per-parameter feature vector fed to the optimiser MLP. Multi-decay EMAs carry the decay along a trailing axis.

init_ema ¶

init_ema(grad: Array, num_decays: int) -> Array

Zero EMA buffer of shape grad.shape + (num_decays,).

update_momentum ¶

update_momentum(momentum: Array, grad: Array, decays: Array = MOMENTUM_DECAYS) -> Array

Multi-decay momentum EMA: m = decay*m + (1-decay)*grad (decay on the last axis).

update_rms ¶

update_rms(rms: Array, grad: Array, decays: Array = MOMENTUM_DECAYS) -> Array

Multi-decay second-moment EMA: rms = decay*rms + (1-decay)*grad**2.

second_moment_normalize ¶

second_moment_normalize(x: Array, axis: int = 0, eps: float = 1e-05) -> Array

Scale features to unit second moment along axis.

Matches mlp_lopt._second_moment_normalizer.

tanh_time_embedding ¶

tanh_time_embedding(iteration: Array) -> Array

Tanh embedding of the iteration over 11 timescales (mlp_lopt._tanh_embedding).

Returns a length-11 vector tanh(iteration / timescale - 1) — a smooth, bounded encoding of training progress shared across all parameters of a tensor.

safe_rsqrt ¶

safe_rsqrt(x: Array) -> Array

Reciprocal square root with a floor (common.safe_rsqrt).

factored_dims ¶

factored_dims(shape: tuple[int, ...]) -> tuple[int, int] | None

Adafactor factoring dims: the two largest axes, or None if rank < 2.

Matches learned_optimizers/common.factored_dims: factored second-moment estimation only applies to tensors of rank >= 2; for those it factors the two largest dimensions. Returns (d1, d0) with d0 the largest axis (reduced for the row estimate) and d1 the second-largest (reduced for the column estimate).

init_adafactor_accum ¶

init_adafactor_accum(param: Array, num_decays: int) -> tuple[Array, Array, Array]

Zeroed factored accumulators (v_row, v_col, v_diag) for one parameter tensor.

Decay is the leading axis. Rank->=2 tensors use (v_row, v_col) (the unused v_diag is an empty placeholder); rank-<2 tensors use a diagonal v_diag (RMSProp-style), mirroring common.factored_rolling.

update_adafactor_accum ¶

update_adafactor_accum(v_row: Array, v_col: Array, v_diag: Array, grad: Array, decays: Array) -> tuple[Array, Array, Array, Array, Array, Array, Array]

Update the factored accumulators and return Adafactor features for one tensor.

Faithful to common.factored_rolling / adafac_mlp_lopt._mod. Returns (new_v_row, new_v_col, new_v_diag, fac_g, row_feat, col_feat, factor) where the four feature arrays carry a trailing decay axis (grad.shape + (num_decays,)): fac_g is the Adafactor-preconditioned gradient, row_feat/col_feat the raw row/column second-moment estimates broadcast back to the tensor shape, and factor the row_factor * col_factor preconditioner (the diagonal rsqrt in the non-factored case).

Coordinatewise learned optimisers (meta-learned update rules).

A :class:LearnedOptimizer carries meta-parameters theta and an opt_fn(theta) that bakes theta into an :class:~opifex.optimization.l2o.optimizers.Optimizer. theta is a plain pytree (the optimiser MLP's :class:flax.nnx state, obtained via nnx.split), so it is directly perturbable/vmappable for evolution-strategies meta-training (PES).

:class:MLPLearnedOptimizer is the per-parameter MLP design of Metz et al. 2020 (arXiv:2009.11243; "LOLv2", learned_optimization/learned_optimizers/mlp_lopt.py): a tiny MLP, shared across all scalar parameters, maps a 19-feature per-parameter vector to a (direction, magnitude) pair, and the update is step = direction * exp(magnitude * exp_mult) * step_mult. The richer Adafactor-feature variant (adafac_mlp_lopt.py) extends the same scaffolding and is added on top of this base.

LearnedOptimizer ¶

Bases: ABC

A meta-learned optimiser: init samples theta; opt_fn(theta) applies it.

init `abstractmethod` ¶

init(key: Array) -> State

Sample the meta-parameters theta (the optimiser MLP state).

opt_fn `abstractmethod` ¶

opt_fn(theta: State) -> Optimizer

Return an :class:Optimizer whose update rule is parameterised by theta.

MLPLOptState ¶

Bases: PyTreeNode

Inner state for learned optimisers: optimisee params, momentum EMAs, iteration.

LearnableSGD ¶

LearnableSGD(initial_learning_rate: float = 0.1)

Bases: LearnedOptimizer

SGD with a single learnable log learning rate (learned_optimizers/base.LearnableSGD).

The simplest learned optimiser: theta is one scalar (log_lr). Useful as a meta-training smoke optimiser and for validating the meta-gradient estimator, since the full-unroll meta-gradient w.r.t. log_lr is analytically differentiable.

init ¶

init(key: Array) -> State

Return theta = the initial log learning rate (deterministic).

opt_fn ¶

opt_fn(theta: State) -> Optimizer

Bake theta (the log-lr) into an SGD optimiser.

MLPLearnedOptimizer ¶

MLPLearnedOptimizer(hidden_size: int = 32, hidden_layers: int = 2, exp_mult: float = 0.001, step_mult: float = 0.001)

Bases: LearnedOptimizer

Per-parameter MLP learned optimiser (LOLv2; Metz et al. 2020).

theta is the MLP's nnx parameter state; opt_fn(theta) applies it coordinatewise with the direction * exp(magnitude * exp_mult) * step_mult update.

init ¶

init(key: Array) -> State

Sample fresh MLP meta-parameters theta from key.

opt_fn ¶

opt_fn(theta: State) -> Optimizer

Bake theta into a coordinatewise :class:Optimizer.

AdafacMLPLOptState ¶

Bases: PyTreeNode

Inner state for the Adafactor-MLP lopt: params, momentum/RMS EMAs, factored accums, step.

mom/rms are multi-decay EMA trees; v_row/v_col/v_diag are the Adafactor factored second-moment accumulators (one tree each, matching the params structure).

AdafacMLPLearnedOptimizer ¶

AdafacMLPLearnedOptimizer(hidden_size: int = 32, hidden_layers: int = 2, exp_mult: float = 0.001, step_mult: float = 0.001)

Bases: LearnedOptimizer

Adafactor-feature per-parameter MLP learned optimiser (adafac_mlp_lopt.py).

Extends :class:MLPLearnedOptimizer with Adafactor-style inputs — multi-decay RMS, m * rsqrt(rms), rsqrt(rms), and factored row/column second-moment features (Metz et al. 2020). theta is the optimiser MLP's nnx state; the update head is the same direction * exp(magnitude * exp_mult) * step_mult rule.

init ¶

init(key: Array) -> State

Sample fresh MLP meta-parameters theta from key.

opt_fn ¶

opt_fn(theta: State) -> Optimizer

Bake theta into a coordinatewise Adafactor-feature :class:Optimizer.

PES meta-training¶

Persistent Evolution Strategies meta-training (Vicol et al. 2021).

Meta-training of learned optimisers via Persistent Evolution Strategies (PES).

Meta-training searches for optimiser meta-parameters theta that minimise the inner task loss accumulated over an unroll, across a distribution of tasks. Back-propagating through a long inner unroll is biased (short truncations) or has exploding/chaotic gradients (long truncations) — Metz et al. 2019 (arXiv:1810.10180). PES (Vicol, Metz & Sohl-Dickstein 2021, arXiv:2112.13835) instead estimates the meta-gradient with antithetic Gaussian perturbations of theta over short truncations, while keeping a persistent accumulator of the perturbations across truncation boundaries so the estimate is unbiased w.r.t. the full-horizon objective. Faithful to learned_optimization/outer_trainers/truncated_pes.py (compute_pes_grad): es_grad = (1 / (2 std**2)) * delta_loss * accumulator.

This implementation runs num_tasks inner problems in parallel (jax.vmap). Each trajectory is started at a random clock offset in [0, total_horizon) (random_initial_iteration_offset in learned_optimization's lopt_truncated_step) and is reset per inner step when its clock reaches total_horizon; the truncation's meta-gradient is split at that reset (has_finished = cumsum(is_done) > 0) so the pre-reset losses attribute to the full accumulator and the post-reset losses only to the new perturbation. Staggering the truncations so the parallel tasks are not phase-aligned is load-bearing (learned_optimization/outer_trainers/truncation_schedule.py): it removes the sawtooth that a synchronous reset would imprint on the meta-loss and lowers the PES gradient variance.

PESState ¶

Bases: PyTreeNode

Persistent PES state, batched over num_tasks inner problems.

task_keys fixes which task each parallel trajectory optimises (so the persistent inner state corresponds to a consistent task); inner_state is the batched inner optimiser state; accumulator is the batched, theta-shaped sum of perturbations since the last horizon reset; inner_step is the per-task inner-step count in the current trajectory (shape (num_tasks,)) — per-task so resets stagger across the parallel trajectories.

pes_gradient_step ¶

pes_gradient_step(learned_optimizer: LearnedOptimizer, task_family: TaskFamily, theta: PyTree, pes_state: PESState, key: Array, *, std: float, trunc_length: int, total_horizon: int) -> tuple[Array, PyTree, PESState]

One PES truncation: return (mean_loss, meta_gradient, new_pes_state).

Faithful to truncated_pes.compute_pes_grad. Antithetic perturbations theta +/- pos are unrolled trunc_length steps from the persistent inner state; the per-step delta-losses are split at the (per-step) horizon reset by has_finished = cumsum(is_done) > 0: losses before the reset attribute to the running accumulator (all perturbations since the last reset), losses after attribute only to the current perturbation pos. The accumulator carries pos forward, or restarts at pos if the trajectory reset during this truncation.

init_pes_state ¶

init_pes_state(learned_optimizer: LearnedOptimizer, task_family: TaskFamily, theta: PyTree, key: Array, *, num_tasks: int, total_horizon: int, trunc_length: int = 1) -> PESState

Build the initial :class:PESState with per-task staggered truncation phases.

Each trajectory starts from a fresh inner state but with a random clock offset in [0, total_horizon) (random_initial_iteration_offset in the reference lopt_truncated_step): because the per-step reset in :func:pes_gradient_step then fires at the staggered clock crossings, the parallel trajectories are never phase-aligned, which removes the sawtooth a synchronous reset would imprint on the meta-loss and lowers the PES variance.

meta_train ¶

meta_train(learned_optimizer: LearnedOptimizer, task_family: TaskFamily, key: Array, *, num_outer_steps: int = 1000, num_tasks: int = 16, trunc_length: int = 20, total_horizon: int = 100, perturbation_std: float = 0.01, meta_learning_rate: float = 0.003) -> tuple[PyTree, Array]

Meta-train learned_optimizer on task_family with PES + outer Adam.

Returns the trained theta and the per-outer-step mean loss curve. PES does not back-propagate through the unroll: theta is updated by the ES estimate fed to Adam.

Baselines and benchmarking¶

optimistix classical baselines and honest learning-curve / speedup-at-target benchmarking.

Real classical optimisation baselines for honest L2O comparison.

A learned optimiser is only interesting if it beats a properly tuned classical optimiser, so the baselines here are genuine: optimistix second-order/line-search minimisers (BFGS/GradientDescent/NonlinearCG) run to convergence, and a tuned first-order optax optimiser (learning-rate swept on the task) for a step-by-step learning-curve comparison. No fabricated baselines or speedups (cf. the deleted _traditional_fallback).

Honest benchmarking follows the L2O literature (Andrychowicz et al. 2016; the VeLO-scaling critique arXiv:2310.18191): compare learning curves and report speedup against a tuned baseline, never an arbitrary one.

optimistix_minimise ¶

optimistix_minimise(task: Task, solver: AbstractMinimiser, start: PyTree, *, max_steps: int = 256, key: Array | None = None) -> tuple[PyTree, Array]

Run a real optimistix minimiser on task from start.

Adapts Task.loss(params, key) to optimistix's fn(y, args) -> scalar (a fixed key, since classical solvers need a deterministic objective). Returns (minimiser, final_loss).

loss_curve ¶

loss_curve(optimizer: Optimizer, task: Task, start: PyTree, *, num_steps: int, key: Array) -> Array

Return the per-step loss curve of an :class:Optimizer on task.

Index 0 is the loss at start; index t is the loss after t update steps.

tuned_optax_baseline ¶

tuned_optax_baseline(task: Task, start: PyTree, learning_rates: Array, *, num_steps: int, key: Array, transformation: Callable[[Array], GradientTransformation] = adam) -> tuple[Array, Array]

Sweep learning_rates and return the best optimiser's (loss_curve, learning_rate).

"Tuned" means the learning rate is chosen by a real sweep (best final loss) — the honest classical first-order baseline for a learning-curve comparison against a learned optimiser.

Honest learned-optimiser benchmarking: learning curves and speedup-at-target.

The primary metric is the loss-vs-step learning curve; the secondary metric is speedup-at-target-loss = (baseline steps to reach the target) / (learned-optimiser steps), censored when a method never reaches the target.

The baseline is tuned with the standard L2O protocol (Andrychowicz et al. 2016): a single learning rate is selected on a tuning batch from the task family (best mean final loss) and then applied unchanged to every held-out task. This is the honest comparison — a learned optimiser's value is precisely that it adapts per-task and per-coordinate without re-tuning, whereas a per-task learning-rate sweep would be an undeployable oracle. Generalisation claims are scoped to in-distribution held-out tasks (cf. the VeLO-scaling critique, arXiv:2310.18191).

steps_to_target ¶

steps_to_target(curve: Array, target: Array) -> Array

First step index at which curve reaches target (inf if never).

speedup_at_target ¶

speedup_at_target(baseline_curve: Array, candidate_curve: Array, target: Array) -> Array

Speedup = baseline-steps / candidate-steps to reach target (censored).

Returns 0 when the candidate never reaches the target, inf when only the baseline fails to — so the value is never fabricated when a method does not converge.

distribution_tuned_lr ¶

distribution_tuned_lr(task_family: TaskFamily, learning_rates: Array, key: Array, *, num_steps: int, num_tune_tasks: int = 16, transformation: Callable[[Array], GradientTransformation] = adam) -> Array

Select one learning rate minimising the mean final loss over a tuning batch.

The standard L2O baseline protocol: tune a single hyperparameter on the task distribution, then apply it unchanged to held-out tasks. Returns the chosen scalar rate.

benchmark_on_held_out_tasks ¶

benchmark_on_held_out_tasks(learned_optimizer: LearnedOptimizer, theta: PyTree, task_family: TaskFamily, key: Array, *, num_tasks: int = 16, num_steps: int = 100, learning_rates: Array = DEFAULT_LR_SWEEP, target_fraction: float = 0.1, transformation: Callable[[Array], GradientTransformation] = adam) -> dict[str, Array]

Meta-test the learned optimiser against a distribution-tuned baseline on held-out tasks.

A single baseline learning rate is tuned on a separate batch from task_family and applied to every held-out task. For each held-out task: run the learned optimiser and the fixed-rate baseline for num_steps, then compute speedup at a per-task target loss (target_fraction of the baseline's initial loss). Returns mean learning curves, the per-task speedups, the median speedup (robust to censored inf/0 entries), and the tuned baseline rate.

L2O Engine¶

High-level orchestrator: meta-train a learned optimiser on a task family, apply it, benchmark it honestly, and persist theta.

High-level learn-to-optimize engine: meta-train, apply, benchmark, persist.

A thin orchestration over the L2O building blocks — it meta-trains a :class:LearnedOptimizer on a :class:TaskFamily with PES, applies the trained optimiser to new tasks, benchmarks it honestly against a tuned classical baseline on held-out tasks, and serialises the meta-learned parameters theta. There is no hidden objective or fabricated speedup: every task carries its own loss and every reported number is measured.

L2OEngine ¶

L2OEngine(learned_optimizer: LearnedOptimizer, task_family: TaskFamily)

Meta-train a learned optimiser on a task family, then apply/benchmark/persist it.

meta_train ¶

meta_train(key: Array, **kwargs: object) -> Array

Meta-train the optimiser (PES); store theta and return the loss curve.

optimize ¶

optimize(task: Task, start: PyTree, *, num_steps: int, key: Array) -> Array

Apply the trained optimiser to task from start; return its loss curve.

benchmark ¶

benchmark(key: Array, **kwargs: object) -> dict[str, Array]

Benchmark the trained optimiser against a tuned baseline on held-out tasks.

save_theta ¶

save_theta(directory: Path) -> None

Serialise the trained theta to directory with Orbax.

load_theta ¶

load_theta(directory: Path, template: PyTree) -> PyTree

Restore theta from directory (template gives the target structure).

Control Systems¶

Differentiable predictive control components for scientific machine learning.

System Identification¶

Neural networks that learn system dynamics from data.

System Identification Networks for Learn-to-Optimize (L2O).

This module implements neural network-based system identification that learns to model dynamical systems with physics constraints, online adaptation, and control integration.

Key Features: - Neural networks for learning system dynamics - Physics-constrained system identification - Online learning and adaptation capabilities - Integration with control policy optimization - Validation on benchmark control systems

PhysicsConstraint `dataclass` ¶

PhysicsConstraint(*, name: str, constraint_type: str, tolerance: float = 0.001, weight: float = 1.0)

Represents a physics constraint for system identification.

This class encapsulates physical laws and constraints that must be enforced during the learning process.

BenchmarkValidationResult `dataclass` ¶

BenchmarkValidationResult(*, benchmark_name: str, metrics: dict[str, float], validation_passed: bool, details: dict[str, Any] | None = None)

Results from benchmark validation.

SystemIdentifier ¶

SystemIdentifier(state_dim: int, input_dim: int, hidden_dim: int = 64, num_layers: int = 3, activation: Callable = gelu, *, rngs: Rngs, dtype: dtype = float32)

Bases: Module

Neural network-based system identification.

This module learns to predict the next state of a dynamical system given the current state and input.

validate_on_benchmark ¶

validate_on_benchmark(benchmark_name: str, test_data: dict[str, Array]) -> BenchmarkValidationResult

Validate system identification on benchmark problem.

Parameters:

Name	Type	Description	Default
`benchmark_name`	`str`	Name of the benchmark	required
`test_data`	`dict[str, Array]`	Test data containing states, inputs, targets	required

Returns:

Type	Description
`BenchmarkValidationResult`	Validation results

integrate_with_l2o_solver ¶

integrate_with_l2o_solver() -> dict[str, Any]

Integration interface with L2O optimization components.

Returns:

Type	Description
`dict[str, Any]`	Integration status and configuration

PhysicsConstrainedSystemID ¶

PhysicsConstrainedSystemID(state_dim: int, input_dim: int, constraints: Sequence[PhysicsConstraint], hidden_dim: int = 64, *, rngs: Rngs, dtype: dtype = float32)

Bases: SystemIdentifier

Physics-constrained system identification.

Extends basic system identification with physics constraints and conservation laws.

compute_energy ¶

compute_energy(state: Array) -> Array

Compute energy of the system state.

Parameters:

Name	Type	Description	Default
`state`	`Array`	System state vector	required

Returns:

Type	Description
`Array`	Scalar energy value

predict_with_constraints ¶

predict_with_constraints(state: Array, input_val: Array) -> dict[str, Any]

Predict next state with constraint checking.

Parameters:

Name	Type	Description	Default
`state`	`Array`	Current state	required
`input_val`	`Array`	Input vector	required

Returns:

Type	Description
`dict[str, Any]`	Dictionary with prediction and constraint violation info

OnlineSystemLearner ¶

OnlineSystemLearner(state_dim: int, input_dim: int, learning_rate: float = 0.001, adaptation_rate: float = 0.95, buffer_size: int = 100, adaptive_lr: bool = False, *, rngs: Rngs, dtype: dtype = float32)

Bases: SystemIdentifier

Online learning system identification.

Adapts the system model in real-time based on new observations.

update_online ¶

update_online(state: Array, input_val: Array, target: Array) -> dict[str, Any]

Update model with new observation.

Parameters:

Name	Type	Description	Default
`state`	`Array`	Current state	required
`input_val`	`Array`	Input that was applied	required
`target`	`Array`	Observed next state	required

Returns:

Type	Description
`dict[str, Any]`	Update results including loss and adaptation metrics

get_memory_info ¶

get_memory_info() -> dict[str, Any]

Get memory management information.

Returns:

Type	Description
`dict[str, Any]`	Memory statistics

ControlIntegratedSystemID ¶

ControlIntegratedSystemID(state_dim: int, input_dim: int, control_dim: int, hidden_dim: int = 64, *, rngs: Rngs, dtype: dtype = float32)

Bases: SystemIdentifier

System identification integrated with control policy learning.

Jointly optimizes system identification and control policy for improved performance.

compute_control_action ¶

compute_control_action(current_state: Array, target_state: Array) -> Array

Compute control action to reach target state.

Parameters:

Name	Type	Description	Default
`current_state`	`Array`	Current system state	required
`target_state`	`Array`	Desired target state	required

Returns:

Type	Description
`Array`	Control action

joint_optimization ¶

joint_optimization(states: Array, targets: Array) -> dict[str, float]

Joint optimization of system ID and control policy.

Parameters:

Name	Type	Description	Default
`states`	`Array`	State trajectory	required
`targets`	`Array`	Target trajectory	required

Returns:

Type	Description
`dict[str, float]`	Optimization losses

simulate_closed_loop ¶

simulate_closed_loop(initial_state: Array, target_state: Array, steps: int) -> dict[str, Array]

Simulate closed-loop system with learned control.

Parameters:

Name	Type	Description	Default
`initial_state`	`Array`	Starting state	required
`target_state`	`Array`	Target state to reach	required
`steps`	`int`	Number of simulation steps	required

Returns:

Type	Description
`dict[str, Array]`	Simulation results

validate_control_benchmark ¶

validate_control_benchmark(benchmark_name: str, initial_state: Array, reference_trajectory: Array, steps: int) -> BenchmarkValidationResult

Validate control performance on benchmark.

Parameters:

Name	Type	Description	Default
`benchmark_name`	`str`	Name of control benchmark	required
`initial_state`	`Array`	Starting state	required
`reference_trajectory`	`Array`	Desired trajectory	required
`steps`	`int`	Number of steps	required

Returns:

Type	Description
`BenchmarkValidationResult`	Validation results

integrate_constraint_learning ¶

integrate_constraint_learning() -> dict[str, Any]

Integration with constraint learning from Version 5.1.

Returns:

Type	Description
`dict[str, Any]`	Constraint satisfaction integration results

SystemDynamicsModel ¶

SystemDynamicsModel(model_type: str, state_dim: int, input_dim: int, hidden_dims: Sequence[int] | None = None, *, rngs: Rngs, dtype: dtype = float32)

Bases: Module

Parameterizable system dynamics model.

Supports both linear and nonlinear system representations.

Model Predictive Control¶

Differentiable MPC frameworks with safety guarantees.

Model Predictive Control (MPC) Framework for Opifex.

Provides differentiable MPC implementation with neural network-based predictive models, constraint handling and projection, real-time control policy optimization, and safety-critical system support.

MPCConfig `dataclass` ¶

MPCConfig(*, horizon: int = 10, control_dim: int = 2, state_dim: int = 4, prediction_steps: int | None = None, objective_weights: dict[str, float] | None = None, max_iterations: int = 50, tolerance: float = 0.0001, time_limit: float = 0.01, learning_rate: float = 0.01)

Configuration for MPC controller.

MPCResult ¶

Bases: NamedTuple

Result from MPC computation.

OptimizationResult ¶

Bases: NamedTuple

Result from optimization.

BatchMPCResult ¶

Bases: NamedTuple

Result from batch MPC computation.

PredictiveModel ¶

PredictiveModel(state_dim: int, control_dim: int, hidden_dims: list[int] | None = None, prediction_horizon: int = 10, model_type: str = 'neural', physics_informed: bool = False, conservation_laws: list[str] | None = None, *, rngs: Rngs)

Bases: Module

Neural network-based predictive model for system dynamics.

predict_step ¶

predict_step(state: ndarray, control: ndarray) -> ndarray

Predict next state given current state and control.

predict_trajectory ¶

predict_trajectory(initial_state: ndarray, control_sequence: ndarray) -> ndarray

Predict state trajectory given control sequence.

ConstraintProjector ¶

ConstraintProjector(state_dim: int, control_dim: int, state_bounds: dict[str, list[float]] | None = None, control_bounds: dict[str, list[float]] | None = None, safety_constraints: bool = False, *, rngs: Rngs)

Bases: Module

Neural network-based constraint projection.

add_custom_constraint ¶

add_custom_constraint(constraint_fn: Callable) -> None

Add custom constraint function.

project_state ¶

project_state(state: ndarray) -> ndarray

Project state to satisfy constraints.

project_control ¶

project_control(control: ndarray) -> ndarray

Project control to satisfy constraints.

ControlBarrier ¶

ControlBarrier(constraint: Callable, alpha: float = 1.0)

Control barrier function for safety.

is_safe_control ¶

is_safe_control(state: ndarray, control: ndarray) -> bool

Check if control is safe given current state.

RealTimeOptimizer ¶

RealTimeOptimizer(max_iterations: int = 50, tolerance: float = 0.0001, learning_rate: float = 0.01, warm_start: bool = True, time_limit: float = 0.01)

Bases: Module

Real-time optimizer for MPC problems.

optimize ¶

optimize(objective: Callable, constraints: Callable | None, initial_guess: ndarray, warm_start_solution: ndarray | None = None) -> OptimizationResult

Optimize objective subject to constraints.

Note: This is not JIT-compatible due to time limits.

Parameters:

Name	Type	Description	Default
`objective`	`Callable`	Objective function to optimize.	required
`constraints`	`Callable \| None`	Constraints to enforce.	required
`initial_guess`	`ndarray`	Initial guess for the solution.	required
`warm_start_solution`	`ndarray \| None`	Solution from previous iteration for warm start.	`None`

Returns:

Name	Type	Description
`OptimizationResult`	`OptimizationResult`	Result of the optimization.

optimize_with_time_limit ¶

optimize_with_time_limit(objective: Callable, constraints: Callable | None, initial_guess: ndarray, warm_start_solution: ndarray | None = None) -> OptimizationResult

Optimize with real-time constraints (not JIT-compatible due to time limits).

This method includes time limit enforcement and therefore cannot be JIT-compiled. Use optimize() for JIT-compatible optimization without time limits.

MPCObjective ¶

MPCObjective(weights: dict[str, float])

MPC objective function.

DifferentiableMPC ¶

DifferentiableMPC(config: MPCConfig, dynamics_model: PredictiveModel | None = None, constraint_projector: ConstraintProjector | None = None)

Bases: Module

Differentiable Model Predictive Control implementation.

set_dynamics ¶

set_dynamics(dynamics_fn: Callable) -> None

Set custom dynamics function.

compute_objective ¶

compute_objective(states: ndarray, controls: ndarray, reference: ndarray) -> Array

Compute MPC objective function.

compute_control ¶

compute_control(current_state: ndarray, reference_trajectory: ndarray) -> MPCResult

Compute optimal control action.

compute_control_batch ¶

compute_control_batch(batch_states: ndarray, batch_references: ndarray) -> BatchMPCResult

Compute control for batch of states.

SafetyCriticalMPC ¶

SafetyCriticalMPC(horizon: int = 10, control_dim: int = 2, state_dim: int = 4, safety_barriers: bool = True, emergency_control: bool = True, backup_policy: bool = True, **kwargs)

Bases: DifferentiableMPC

Safety-critical MPC with emergency control and backup policies.

add_barrier ¶

add_barrier(barrier: ControlBarrier) -> None

Add control barrier function.

compute_safe_control ¶

compute_safe_control(current_state: ndarray, reference_trajectory: ndarray) -> MPCResult

Compute safe control action with emergency and backup policies.

RecedingHorizonController ¶

RecedingHorizonController(mpc_horizon: int = 10, control_horizon: int = 5, state_dim: int = 4, control_dim: int = 2, sampling_time: float = 0.1, safety_critical: bool = False)

Bases: Module

Receding horizon controller implementation.

compute_control ¶

compute_control(current_state: ndarray, reference_trajectory: ndarray) -> MPCResult

Compute control using receding horizon.

simulate_tracking ¶

simulate_tracking(initial_state: ndarray, reference_trajectory: ndarray, simulation_steps: int | None = None) -> ndarray

Simulate reference tracking.

Module Overview¶

The optimization module is organized into several key components:

Core Components¶

meta_optimization/: Meta-optimization framework with L2O algorithms (modular package)
production.py: Production optimization (adaptive JIT, GPU memory planning)
scientific_integration.py: Physics-aware optimization integration

L2O Submodule (`l2o/`)¶

core.py: Task/TaskFamily (objective-carrying) and the Optimizer interface
optimizers.py: Optimizer ABC + OptaxOptimizer (hand-designed baseline family)
tasks.py: QuadraticTaskFamily and the MLPTaskFamily showcase task
features.py: per-parameter input features (momentum/RMS, time embedding)
learned.py: LearnedOptimizer ABC, MLPLearnedOptimizer, LearnableSGD
meta_train.py: Persistent Evolution Strategies (PES) meta-training
baselines.py: optimistix classical baselines and tuned-optax baselines
benchmark.py: honest learning-curve and speedup-at-target benchmarking
engine.py: high-level L2OEngine orchestrator

Control Submodule (`control/`)¶

system_id.py: System identification networks
mpc.py: Model predictive control frameworks

Key Features¶

Meta-Optimization Features¶

Learn-to-Optimize (L2O) algorithms with meta-learned update rules
Adaptive learning rate scheduling
Warm-starting strategies for related problems

Production Optimization Features¶

Hybrid performance platform with adaptive JIT
Intelligent GPU memory management
Physics-aware scientific validation of optimized models

Control Systems Features¶

Differentiable model predictive control
Physics-constrained system identification
Safety-critical control with barrier functions
Real-time optimization capabilities

Scientific Integration Features¶

Physics-informed optimization
Conservation law enforcement
Numerical validation and stability checks
Domain-specific profiling and benchmarking

Usage Examples¶

Basic Meta-Optimization¶

from opifex.optimization.meta_optimization import LearnToOptimize, MetaOptimizerConfig

config = MetaOptimizerConfig(
    meta_learning_rate=1e-4,
    adaptation_steps=5,
    warm_start_strategy="previous_params"
)

l2o = LearnToOptimize(config=config, rngs=nnx.Rngs(42))
optimized_params = l2o.optimize(params, objective_fn, num_steps=1000)

Production Optimization¶

from opifex.optimization.production import HybridPerformancePlatform, WorkloadProfile

platform = HybridPerformancePlatform()

workload = WorkloadProfile(
    batch_size=32,
    sequence_length=128,
    memory_footprint=2.0,
    compute_intensity=8.0,
    latency_requirement=10.0,
    throughput_requirement=100.0,
    model_complexity="medium",
)

optimized = platform.optimize_for_production(model, workload)

Control System¶

from opifex.optimization.control import DifferentiableMPC, SystemIdentifier

# Learn system dynamics
system_id = SystemIdentifier(model=dynamics_model)
trained_model = system_id.fit(state_data, input_data)

# Create MPC controller
mpc = DifferentiableMPC(system_model=trained_model, config=mpc_config)
control_action = mpc.solve(current_state, reference_trajectory)

Performance Characteristics¶

L2O Speedup: meta-learned optimizers accelerate convergence on learned problem families
Meta-Optimization: faster convergence on related problems via warm-starting
Production Optimization: adaptive JIT kernel fusion with measured speedups
Memory Efficiency: pool-based GPU memory planning for co-located models

Integration¶

The optimization module integrates seamlessly with:

Training: Meta-optimization for training workflows
Neural Networks: Compatible with all neural architectures
Physics: Physics-informed optimization constraints
Deployment: Production-ready optimization systems

Second-Order Optimization¶

Curvature-based optimization methods including L-BFGS and hybrid optimizers.

Configuration Classes¶

opifex.optimization.second_order.config ¶

Configuration classes for second-order optimization methods.

This module provides unified configuration dataclasses for all second-order optimization methods supported by the Opifex framework.

Design Principles

All configs are frozen dataclasses (immutable)
Validation happens at construction time via post_init
Sensible defaults based on literature recommendations
Clear separation between method-specific and shared configs

References

Survey: arXiv:2601.10222v1 Section 7
L-BFGS memory size: typically 3-20 (Liu & Nocedal, 1989)
Hybrid switching: Section 7.4 of the survey

LBFGSConfig `dataclass` ¶

LBFGSConfig(memory_size: int = 10, scale_init_precond: bool = True, linesearch: LinesearchType = ZOOM, max_linesearch_steps: int = 20, max_iterations: int = 100, tolerance: float = 1e-06)

Configuration for L-BFGS optimizer.

L-BFGS (Limited-memory BFGS) approximates the inverse Hessian using a limited history of gradient differences. This makes it suitable for large-scale optimization where storing the full Hessian is infeasible.

Attributes:

Name	Type	Description
`memory_size`	`int`	Number of gradient pairs to store (typically 3-20)
`scale_init_precond`	`bool`	Whether to scale initial preconditioner
`linesearch`	`LinesearchType`	Line search algorithm to use
`max_linesearch_steps`	`int`	Maximum steps for line search
`max_iterations`	`int`	Maximum L-BFGS iterations
`tolerance`	`float`	Convergence tolerance

References

Liu & Nocedal (1989): On the limited memory BFGS method
optax.lbfgs documentation

GaussNewtonConfig `dataclass` ¶

GaussNewtonConfig(damping_factor: float = 0.001, damping_increase_factor: float = 10.0, damping_decrease_factor: float = 0.1, min_damping: float = 1e-10, max_damping: float = 10000000000.0, max_iterations: int = 100, rtol: float = 1e-06, atol: float = 1e-06)

Configuration for Gauss-Newton and Levenberg-Marquardt solvers.

Gauss-Newton is effective for nonlinear least-squares problems where the residual Jacobian can be computed efficiently. Levenberg-Marquardt adds damping for improved robustness.

Attributes:

Name	Type	Description
`damping_factor`	`float`	Initial damping factor (λ) for LM
`damping_increase_factor`	`float`	Factor to increase damping on failure (> 1)
`damping_decrease_factor`	`float`	Factor to decrease damping on success (< 1)
`min_damping`	`float`	Minimum allowed damping value
`max_damping`	`float`	Maximum allowed damping value
`max_iterations`	`int`	Maximum solver iterations
`rtol`	`float`	Relative tolerance for convergence
`atol`	`float`	Absolute tolerance for convergence

References

optimistix.LevenbergMarquardt documentation
Survey Section 7.3

HybridOptimizerConfig `dataclass` ¶

HybridOptimizerConfig(first_order_steps: int = 1000, switch_criterion: SwitchCriterion = LOSS_VARIANCE, loss_variance_threshold: float = 0.0001, loss_history_window: int = 50, gradient_norm_threshold: float = 0.001, relative_improvement_threshold: float = 0.0001, adam_learning_rate: float = 0.001, adam_b1: float = 0.9, adam_b2: float = 0.999, lbfgs_config: LBFGSConfig = LBFGSConfig())

Configuration for hybrid Adam→L-BFGS optimizer.

This optimizer starts with Adam for initial exploration and switches to L-BFGS for efficient convergence once the loss landscape becomes smooth. This follows recommendations from Survey Section 7.4.

The switch can be triggered by various criteria

EPOCH: Switch after fixed number of steps
LOSS_VARIANCE: Switch when loss variance drops below threshold
GRADIENT_NORM: Switch when gradient norm drops below threshold
RELATIVE_IMPROVEMENT: Switch when relative improvement slows

Attributes:

Name	Type	Description
`first_order_steps`	`int`	Steps to run Adam before considering switch
`switch_criterion`	`SwitchCriterion`	Criterion for switching to L-BFGS
`loss_variance_threshold`	`float`	Threshold for loss variance criterion
`loss_history_window`	`int`	Window size for computing loss statistics
`gradient_norm_threshold`	`float`	Threshold for gradient norm criterion
`relative_improvement_threshold`	`float`	Threshold for relative improvement
`adam_learning_rate`	`float`	Learning rate for Adam phase
`adam_b1`	`float`	Adam beta1 parameter
`adam_b2`	`float`	Adam beta2 parameter
`lbfgs_config`	`LBFGSConfig`	Configuration for L-BFGS phase

References

Survey Section 7.4: "L-BFGS is more effective in later stages when loss varies smoothly"

L-BFGS and Gauss-Newton Wrappers¶

opifex.optimization.second_order.wrappers ¶

Wrappers for external second-order optimization libraries.

This module provides thin wrappers around optax and optimistix to create second-order optimizers with our unified configuration interface.

Design Philosophy

Wrap existing robust implementations (optax, optimistix)
Don't reinvent the wheel
Provide consistent interface through our config classes

References

optax.lbfgs: Pure JAX L-BFGS with line search
optimistix: Gauss-Newton, Levenberg-Marquardt, BFGS

create_lbfgs_optimizer ¶

create_lbfgs_optimizer(config: LBFGSConfig | None = None) -> GradientTransformation

Create L-BFGS optimizer using optax.

L-BFGS is a quasi-Newton method that approximates the inverse Hessian using a limited history of gradient differences. This is the recommended second-order optimizer for large-scale optimization.

Parameters:

Name	Type	Description	Default
`config`	`LBFGSConfig \| None`	L-BFGS configuration. Uses defaults if None.	`None`

Returns:

Type	Description
`GradientTransformation`	Optax L-BFGS gradient transformation.

Example

config = LBFGSConfig(memory_size=20) optimizer = create_lbfgs_optimizer(config)

Use with optax training loop¶

create_gauss_newton_solver ¶

create_gauss_newton_solver(config: GaussNewtonConfig | None = None) -> AbstractLeastSquaresSolver

Create Gauss-Newton solver using optimistix.

Gauss-Newton is effective for nonlinear least-squares problems. Note that this creates a solver for root-finding/minimization, not a gradient transformation like L-BFGS.

Parameters:

Name	Type	Description	Default
`config`	`GaussNewtonConfig \| None`	Gauss-Newton configuration. Uses defaults if None.	`None`

Returns:

Type	Description
`AbstractLeastSquaresSolver`	Optimistix Gauss-Newton solver.

Example

solver = create_gauss_newton_solver()

Use with optimistix.least_squares¶

Hybrid Adam → L-BFGS Optimizer¶

opifex.optimization.second_order.hybrid_optimizer ¶

Hybrid Adam→L-BFGS optimizer for physics-informed training.

This module implements a hybrid optimization strategy that starts with Adam for initial exploration and switches to L-BFGS for efficient convergence once the loss landscape becomes smooth.

Design Rationale (from Survey Section 7.4): "L-BFGS is more effective in later stages when loss varies smoothly."

The hybrid approach combines

Adam's robustness in noisy, high-curvature early optimization
L-BFGS's superior convergence in smooth regions near optima

Key Features

Multiple switching criteria (epoch, loss variance, gradient norm)
Loss history tracking for variance-based switching
Full JAX/JIT compatibility
Works with FLAX NNX models

References

Survey: arXiv:2601.10222v1 Section 7.4

HybridOptimizer `dataclass` ¶

HybridOptimizer(config: HybridOptimizerConfig)

Hybrid Adam→L-BFGS optimizer.

This optimizer starts with Adam and switches to L-BFGS based on configurable criteria. The transition is designed to leverage Adam's robustness in early training and L-BFGS's efficiency for final convergence.

Attributes:

Name	Type	Description
`config`	`HybridOptimizerConfig`	Hybrid optimizer configuration
`adam`	`GradientTransformation`	Adam optimizer instance
`lbfgs`	`GradientTransformationExtraArgs`	L-BFGS optimizer instance

Example

config = HybridOptimizerConfig(first_order_steps=1000) optimizer = HybridOptimizer(config) state = optimizer.init(params)

Training loop¶

for step in range(num_steps): ... loss, grads = loss_and_grad_fn(params) ... updates, state = optimizer.update(grads, state, params, loss=loss) ... params = optax.apply_updates(params, updates)

is_using_lbfgs `property` ¶

is_using_lbfgs: bool

Check if optimizer is currently using L-BFGS.

Note: This is a convenience property. For actual state, check the HybridOptimizerState.using_lbfgs field.

init ¶

init(params: PyTree) -> HybridOptimizerState

Initialize optimizer state.

Parameters:

Name	Type	Description	Default
`params`	`PyTree`	Model parameters (PyTree)	required

Returns:

Type	Description
`HybridOptimizerState`	Initial optimizer state

update ¶

update(grads: PyTree, state: HybridOptimizerState, params: PyTree, *, loss: Float[Array, ''] | None = None, value: Float[Array, ''] | None = None, grad: PyTree | None = None, value_fn: Callable[[PyTree], Float[Array, '']] | None = None) -> tuple[PyTree, HybridOptimizerState]

Compute parameter updates.

This method handles the switching logic and delegates to either Adam or L-BFGS depending on the current state.

Parameters:

Name	Type	Description	Default
`grads`	`PyTree`	Parameter gradients	required
`state`	`HybridOptimizerState`	Current optimizer state	required
`params`	`PyTree`	Current parameters (needed for L-BFGS)	required
`loss`	`Float[Array, ''] \| None`	Current loss value (for variance-based switching)	`None`
`value`	`Float[Array, ''] \| None`	Alias for loss (for optax L-BFGS compatibility)	`None`
`grad`	`PyTree \| None`	Alias for grads (for optax L-BFGS compatibility)	`None`
`value_fn`	`Callable[[PyTree], Float[Array, '']] \| None`	Loss function (needed for L-BFGS line search)	`None`

Returns:

Type	Description
`tuple[PyTree, HybridOptimizerState]`	Tuple of (updates, new_state)

NNX Integration¶

opifex.optimization.second_order.nnx_integration ¶

FLAX NNX integration for second-order optimizers.

This module provides wrapper classes that make it easy to use second-order optimizers (L-BFGS, hybrid Adam→L-BFGS) with FLAX NNX models.

Design Philosophy

Hide the complexity of nnx.split/merge from users
Provide familiar step() interface similar to nnx.Optimizer
Support both pure L-BFGS and hybrid optimization strategies

Key Classes

NNXSecondOrderOptimizer: L-BFGS optimizer for NNX models
NNXHybridOptimizer: Hybrid Adam→L-BFGS for NNX models
create_nnx_lbfgs_optimizer: Factory function for L-BFGS

References

FLAX NNX documentation: https://flax.readthedocs.io/en/latest/nnx/
optax L-BFGS requires functional API with value_and_grad_from_state

create_nnx_lbfgs_optimizer ¶

create_nnx_lbfgs_optimizer(model: Module, config: LBFGSConfig | None = None) -> NNXSecondOrderOptimizer

Create L-BFGS optimizer for NNX model.

Factory function that creates an NNXSecondOrderOptimizer configured with L-BFGS.

Parameters:

Name	Type	Description	Default
`model`	`Module`	FLAX NNX model to optimize	required
`config`	`LBFGSConfig \| None`	L-BFGS configuration. Uses defaults if None.	`None`

Returns:

Type	Description
`NNXSecondOrderOptimizer`	Configured NNXSecondOrderOptimizer instance.

Example

model = MyModel(rngs=nnx.Rngs(0)) optimizer = create_nnx_lbfgs_optimizer(model) for _ in range(100): ... loss = optimizer.step(loss_fn)

For detailed algorithms and best practices, see the Second-Order Optimization Guide.

Optimization API Reference¶

Overview¶

Core Optimization Components¶

Meta-Optimization Framework¶

MetaOptimizerConfig dataclass ¶

MetaOptimizer ¶

init_optimizer_state ¶

step ¶

store_optimization_result ¶

get_warm_start_params ¶

quantum_step ¶

PerformanceMonitor ¶

update_metrics ¶

get_metric_history ¶

check_convergence ¶

get_performance_analytics ¶

get_quantum_analytics ¶

LearnToOptimize ¶

compute_update ¶

compute_meta_gradients ¶

compute_adaptive_update ¶

compute_quantum_update ¶

AdaptiveLearningRateScheduler ¶

get_learning_rate ¶

adapt_from_performance ¶

adapt_from_quantum_metrics ¶

WarmStartingStrategy ¶

get_warm_start_params ¶

adapt_optimizer_state ¶

get_molecular_warm_start ¶

Production Optimization¶

CallableModule ¶

OptimizationStrategy ¶

WorkloadProfile dataclass ¶

PerformanceMetrics dataclass ¶

OptimizedModel dataclass ¶

AdaptiveJAXOptimizer ¶

analyze_workload_patterns ¶

apply_aggressive_kernel_fusion ¶

apply_memory_optimization ¶

apply_latency_optimization ¶

apply_balanced_optimization ¶

benchmark_model_performance ¶

optimize_neural_operator ¶

IntelligentGPUMemoryManager ¶

select_memory_pool ¶

estimate_model_memory_usage ¶

optimize_multi_model_allocation ¶

HybridPerformancePlatform ¶

optimize_for_production ¶

get_model_input_features ¶

Scientific Computing Integration¶

PhysicsDomain ¶

PhysicsMetrics dataclass ¶

NumericalValidationResult dataclass ¶

ConservationCheckResult dataclass ¶

ScientificBenchmarkResult dataclass ¶

PhysicsProfilerProtocol ¶

profile_physics_metrics ¶

validate_domain_constraints ¶

PhysicsProfiler ¶

profile_physics_metrics ¶

NumericalValidator ¶

validate_numerical_precision ¶

check_conservation_law ¶

ScientificBenchmarkValidator ¶

validate_benchmark ¶

validate_multiple_benchmarks ¶

ScientificComputingIntegrator ¶

comprehensive_scientific_validation ¶

optimize_for_scientific_accuracy ¶

translation_invariance_error ¶

Learn-to-Optimize (L2O) Algorithms¶

Tasks and the optimiser interface¶

Task ¶

init abstractmethod ¶

loss abstractmethod ¶

normalizer ¶

loss_and_grad ¶

TaskFamily ¶

MetaOptimizerConfig `dataclass` ¶

WorkloadProfile `dataclass` ¶

PerformanceMetrics `dataclass` ¶

OptimizedModel `dataclass` ¶

PhysicsMetrics `dataclass` ¶

NumericalValidationResult `dataclass` ¶

ConservationCheckResult `dataclass` ¶

ScientificBenchmarkResult `dataclass` ¶

init `abstractmethod` ¶

loss `abstractmethod` ¶

sample `abstractmethod` ¶

init `abstractmethod` ¶

update `abstractmethod` ¶

get_params `abstractmethod` ¶

QuadraticTask `dataclass` ¶

QuadraticTaskFamily `dataclass` ¶

MLPTask `dataclass` ¶

MLPTaskFamily `dataclass` ¶

layer_sizes `property` ¶

init `abstractmethod` ¶

opt_fn `abstractmethod` ¶