diff --git a/src/lib.rs b/src/lib.rs index a1fbdf7..32c685a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -145,17 +145,16 @@ pub use mono::MonoAD; pub use mono::{MonoAD2FR, MonoAD2RF, MonoAD2RR}; pub use multi::{ - builder::GraphBuilder, multi_ad_fr::MultiAD2FR, multi_ad_rf::MultiAD2RF, - multi_ad_rr::MultiAD2RR, AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, - BackendKind, BackendRejectionReason, BackendSupportReport, BatchGradients, - BatchGradientsBuffer, BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, - CompiledGraphMetadata, CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, - DeviceBufferHandle, DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode, + AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, BackendKind, + BackendRejectionReason, BackendSupportReport, BatchGradients, BatchGradientsBuffer, + BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, + CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, DeviceBufferHandle, + DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode, DeviceExecutionTrace, DeviceMemoryLocation, DeviceTransferKind, DeviceTransferPlan, DeviceTransferPolicy, DomainPolicy, ExecutionBackend, ExprGraph, ExprNode, FlatInstruction, - GpuBackendBoundary, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, GraphStats, - Instruction, MockDeviceBackend, MultiAD, NodeId, OpCode, ScalarBackend, SimdBackend, Tape, - TapeWorkspace, UNUSED_NODE_ID, + GpuBackendBoundary, GradientCheckEntry, GradientCheckReport, Graph, GraphBuilder, GraphNode, + GraphStats, Instruction, MockDeviceBackend, MultiAD, MultiAD2FR, MultiAD2RF, MultiAD2RR, + NodeId, OpCode, ScalarBackend, SimdBackend, Tape, TapeWorkspace, UNUSED_NODE_ID, }; #[cfg(feature = "backend-wgpu")] pub use multi::{ diff --git a/src/mono.rs b/src/mono.rs deleted file mode 100644 index 08dd72f..0000000 --- a/src/mono.rs +++ /dev/null @@ -1,31 +0,0 @@ -pub mod types; - -#[cfg(test)] -mod tests; - -#[cfg(test)] -mod tests_ho; - -pub mod mono_ad; -pub use mono_ad::MonoAD; - -pub mod mono_ad_rr; -pub use mono_ad_rr::MonoAD2RR; - -pub mod mono_ad_fr; -mod mono_hessian_common; -pub use mono_ad_fr::MonoAD2FR; - -pub mod mono_ad_rf; -pub use mono_ad_rf::MonoAD2RF; - -mod mono_fn; -// Re-export trait for library extension - users can implement custom mono functions -#[allow(unused_imports)] // May not be used internally, but part of public API -pub use mono_fn::MonoFn; - -// Example implementation - not part of public API -mod mf1; -mod mf2; -mod mf3; -mod mf4; diff --git a/src/mono/examples.rs b/src/mono/examples.rs new file mode 100644 index 0000000..93b3e80 --- /dev/null +++ b/src/mono/examples.rs @@ -0,0 +1,101 @@ +//! Test-only example implementations of [`MonoFn`] for library tests. +//! +//! These are production-type examples used internally by the test suite and +//! are gated behind `#[cfg(test)]` to keep them out of release builds. + +#[cfg(test)] +use super::func::{GraphType, MonoFn}; +#[cfg(test)] +use crate::mono_ops; + +/// f(x) = exp(sin(sin(x))) +#[cfg(test)] +pub struct MF1(pub f64); + +#[cfg(test)] +impl MonoFn for MF1 { + fn input(&self) -> f64 { + self.0 + } + + fn graph(&self) -> &'static GraphType { + &mono_ops![sin, sin, exp] + } + + fn expected_value(&self) -> f64 { + (self.0.sin().sin()).exp() + } + + fn expected_gradient(&self) -> f64 { + (self.0.sin().sin()).exp() * self.0.sin().cos() * self.0.cos() + } +} + +/// f(x) = -x +#[cfg(test)] +pub struct MF2(pub f64); + +#[cfg(test)] +impl MonoFn for MF2 { + fn input(&self) -> f64 { + self.0 + } + + fn graph(&self) -> &'static GraphType { + &mono_ops![neg] + } + + fn expected_value(&self) -> f64 { + -self.0 + } + + fn expected_gradient(&self) -> f64 { + -1.0 + } +} + +/// f(x) = sin(-x) +#[cfg(test)] +pub struct MF3(pub f64); + +#[cfg(test)] +impl MonoFn for MF3 { + fn input(&self) -> f64 { + self.0 + } + + fn graph(&self) -> &'static GraphType { + &mono_ops![neg, sin] + } + + fn expected_value(&self) -> f64 { + (-self.0).sin() + } + + fn expected_gradient(&self) -> f64 { + -((-self.0).cos()) + } +} + +/// f(x) = -sin(x) +#[cfg(test)] +pub struct MF4(pub f64); + +#[cfg(test)] +impl MonoFn for MF4 { + fn input(&self) -> f64 { + self.0 + } + + fn graph(&self) -> &'static GraphType { + &mono_ops![sin, neg] + } + + fn expected_value(&self) -> f64 { + -(self.0.sin()) + } + + fn expected_gradient(&self) -> f64 { + -(self.0.cos()) + } +} diff --git a/src/mono/mono_ad.rs b/src/mono/first_order.rs similarity index 100% rename from src/mono/mono_ad.rs rename to src/mono/first_order.rs diff --git a/src/mono/mono_fn.rs b/src/mono/func.rs similarity index 99% rename from src/mono/mono_fn.rs rename to src/mono/func.rs index b48de14..2c4f9c3 100644 --- a/src/mono/mono_fn.rs +++ b/src/mono/func.rs @@ -1,4 +1,4 @@ -pub use super::mono_ad::MonoAD; +pub use super::first_order::MonoAD; pub use super::types::BackwardResultBox; /// Type alias for a graph of mono operations (slice of MonoAD) diff --git a/src/mono/mf1.rs b/src/mono/mf1.rs deleted file mode 100644 index c15ee7e..0000000 --- a/src/mono/mf1.rs +++ /dev/null @@ -1,26 +0,0 @@ -#[cfg(test)] -use super::mono_fn::{GraphType, MonoFn}; -#[cfg(test)] -use crate::mono_ops; - -#[cfg(test)] -pub struct MF1(pub f64); - -#[cfg(test)] -impl MonoFn for MF1 { - fn input(&self) -> f64 { - self.0 - } - - fn graph(&self) -> &'static GraphType { - &mono_ops![sin, sin, exp] - } - - fn expected_value(&self) -> f64 { - (self.0.sin().sin()).exp() - } - - fn expected_gradient(&self) -> f64 { - (self.0.sin().sin()).exp() * self.0.sin().cos() * self.0.cos() - } -} diff --git a/src/mono/mf2.rs b/src/mono/mf2.rs deleted file mode 100644 index ba8af5f..0000000 --- a/src/mono/mf2.rs +++ /dev/null @@ -1,26 +0,0 @@ -#[cfg(test)] -use super::mono_fn::{GraphType, MonoFn}; -#[cfg(test)] -use crate::mono_ops; - -#[cfg(test)] -pub struct MF2(pub f64); - -#[cfg(test)] -impl MonoFn for MF2 { - fn input(&self) -> f64 { - self.0 - } - - fn graph(&self) -> &'static GraphType { - &mono_ops![neg] - } - - fn expected_value(&self) -> f64 { - -self.0 - } - - fn expected_gradient(&self) -> f64 { - -1.0 - } -} diff --git a/src/mono/mf3.rs b/src/mono/mf3.rs deleted file mode 100644 index 9cc51a9..0000000 --- a/src/mono/mf3.rs +++ /dev/null @@ -1,26 +0,0 @@ -#[cfg(test)] -use super::mono_fn::{GraphType, MonoFn}; -#[cfg(test)] -use crate::mono_ops; - -#[cfg(test)] -pub struct MF3(pub f64); - -#[cfg(test)] -impl MonoFn for MF3 { - fn input(&self) -> f64 { - self.0 - } - - fn graph(&self) -> &'static GraphType { - &mono_ops![neg, sin] - } - - fn expected_value(&self) -> f64 { - (-self.0).sin() - } - - fn expected_gradient(&self) -> f64 { - -((-self.0).cos()) - } -} diff --git a/src/mono/mf4.rs b/src/mono/mf4.rs deleted file mode 100644 index 3e4f6ec..0000000 --- a/src/mono/mf4.rs +++ /dev/null @@ -1,26 +0,0 @@ -#[cfg(test)] -use super::mono_fn::{GraphType, MonoFn}; -#[cfg(test)] -use crate::mono_ops; - -#[cfg(test)] -pub struct MF4(pub f64); - -#[cfg(test)] -impl MonoFn for MF4 { - fn input(&self) -> f64 { - self.0 - } - - fn graph(&self) -> &'static GraphType { - &mono_ops![sin, neg] - } - - fn expected_value(&self) -> f64 { - -(self.0.sin()) - } - - fn expected_gradient(&self) -> f64 { - -(self.0.cos()) - } -} diff --git a/src/mono/mod.rs b/src/mono/mod.rs new file mode 100644 index 0000000..90bbf56 --- /dev/null +++ b/src/mono/mod.rs @@ -0,0 +1,28 @@ +//! Single-variable automatic differentiation. +//! +//! This module provides functionality for computing derivatives of +//! single-variable functions using reverse-mode differentiation. + +pub mod types; + +#[cfg(test)] +mod tests; + +#[cfg(test)] +mod tests_ho; + +pub mod first_order; +pub use first_order::MonoAD; + +pub mod second_order; +pub use second_order::fr::MonoAD2FR; +pub use second_order::rf::MonoAD2RF; +pub use second_order::rr::MonoAD2RR; + +pub mod func; +// Re-export trait for library extension - users can implement custom mono functions +#[allow(unused_imports)] +pub use func::MonoFn; + +#[cfg(test)] +mod examples; diff --git a/src/mono/mono_hessian_common.rs b/src/mono/second_order/common.rs similarity index 99% rename from src/mono/mono_hessian_common.rs rename to src/mono/second_order/common.rs index 34b7178..7b928f1 100644 --- a/src/mono/mono_hessian_common.rs +++ b/src/mono/second_order/common.rs @@ -30,7 +30,7 @@ use crate::{AutodiffError, Result}; -use super::types::*; +use crate::mono::types::*; /// Operation kind for mono-variable second-order AD (shared by FR and RF). /// diff --git a/src/mono/mono_ad_fr.rs b/src/mono/second_order/fr.rs similarity index 93% rename from src/mono/mono_ad_fr.rs rename to src/mono/second_order/fr.rs index e65d3e2..f462ef0 100644 --- a/src/mono/mono_ad_fr.rs +++ b/src/mono/second_order/fr.rs @@ -12,7 +12,7 @@ //! # Supported Operations //! //! This type supports a subset of operations compared to [`crate::MonoAD`]: -//! `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. See [`super::mono_hessian_common`] for details. +//! `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. See [`super::common`] for details. //! //! # Mathematical Foundation //! @@ -65,9 +65,9 @@ use crate::Result; -use super::mono_ad_rr::MonoAD2RR; -use super::mono_hessian_common::{self, MonoHessianOpKind}; -use super::types::*; +use super::common::{self, MonoHessianOpKind}; +use super::rr::MonoAD2RR; +use crate::mono::types::*; /// Single-variable automatic differentiation operations for Forward-over-Reverse Hessian computation. /// @@ -103,13 +103,13 @@ impl MonoAD2FR { /// Compute forward pass only. pub fn compute(exprs: &[MonoAD2FR], x: f64) -> f64 { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_forward(&ops, x) + common::compute_forward(&ops, x) } /// Compute forward pass with opt-in checked-domain validation. pub fn compute_checked(exprs: &[MonoAD2FR], x: f64) -> Result { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_forward_checked(&ops, x) + common::compute_forward_checked(&ops, x) } /// Compute forward pass and return gradient function using reverse-mode. @@ -120,7 +120,7 @@ impl MonoAD2FR { /// Compute forward pass and gradient function with checked-domain validation. pub fn compute_grad_checked(exprs: &[MonoAD2FR], x: f64) -> Result { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_grad_generic_checked::>(&ops, x) + common::compute_grad_generic_checked::>(&ops, x) } fn compute_grad_generic(exprs: &[MonoAD2FR], x: f64) -> (f64, W) @@ -128,7 +128,7 @@ impl MonoAD2FR { W: From> + std::ops::Deref + 'static, { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_grad_generic::(&ops, x) + common::compute_grad_generic::(&ops, x) } /// Compute exact Hessian using Forward-over-Reverse mode. @@ -171,7 +171,7 @@ impl MonoAD2FR { // Single operation: direct dual-number evaluation if let [op] = exprs { - return mono_hessian_common::compute_single_op_hessian((*op).into(), x) + return common::compute_single_op_hessian((*op).into(), x) .expect("all single ops supported"); } diff --git a/src/mono/second_order/mod.rs b/src/mono/second_order/mod.rs new file mode 100644 index 0000000..4fc7538 --- /dev/null +++ b/src/mono/second_order/mod.rs @@ -0,0 +1,6 @@ +//! Exact second-order (Hessian) computation for single-variable functions. + +pub(crate) mod common; +pub mod fr; +pub mod rf; +pub mod rr; diff --git a/src/mono/mono_ad_rf.rs b/src/mono/second_order/rf.rs similarity index 94% rename from src/mono/mono_ad_rf.rs rename to src/mono/second_order/rf.rs index f119b99..ea5314f 100644 --- a/src/mono/mono_ad_rf.rs +++ b/src/mono/second_order/rf.rs @@ -11,7 +11,7 @@ //! # Supported Operations //! //! This type supports a subset of operations compared to [`crate::MonoAD`]: -//! `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. See [`super::mono_hessian_common`] for details. +//! `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. See [`super::common`] for details. //! //! For single-variable functions, RF and FR are mathematically equivalent (both use dual //! numbers), but they differ in conceptual organization: @@ -73,7 +73,7 @@ //! - **RF**: Differentiate function (forward) for each output (reverse) //! //! For MonoAD (single variable), this implementation is shared with FR via -//! [`super::mono_hessian_common`]. +//! [`super::common`]. //! //! # Algorithm //! @@ -117,9 +117,9 @@ use crate::Result; -use super::mono_ad_rr::MonoAD2RR; -use super::mono_hessian_common::{self, MonoHessianOpKind}; -use super::types::*; +use super::common::{self, MonoHessianOpKind}; +use super::rr::MonoAD2RR; +use crate::mono::types::*; /// Single-variable automatic differentiation operations for Reverse-over-Forward Hessian computation. /// @@ -155,13 +155,13 @@ impl MonoAD2RF { /// Compute forward pass only. pub fn compute(exprs: &[MonoAD2RF], x: f64) -> f64 { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_forward(&ops, x) + common::compute_forward(&ops, x) } /// Compute forward pass with opt-in checked-domain validation. pub fn compute_checked(exprs: &[MonoAD2RF], x: f64) -> Result { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_forward_checked(&ops, x) + common::compute_forward_checked(&ops, x) } /// Compute forward pass and return gradient function using reverse-mode. @@ -172,7 +172,7 @@ impl MonoAD2RF { /// Compute forward pass and gradient function with checked-domain validation. pub fn compute_grad_checked(exprs: &[MonoAD2RF], x: f64) -> Result { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_grad_generic_checked::>(&ops, x) + common::compute_grad_generic_checked::>(&ops, x) } fn compute_grad_generic(exprs: &[MonoAD2RF], x: f64) -> (f64, W) @@ -180,7 +180,7 @@ impl MonoAD2RF { W: From> + std::ops::Deref + 'static, { let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); - mono_hessian_common::compute_grad_generic::(&ops, x) + common::compute_grad_generic::(&ops, x) } /// Compute exact Hessian using Reverse-over-Forward mode. @@ -227,7 +227,7 @@ impl MonoAD2RF { // Single operation: direct dual-number evaluation if let [op] = exprs { - return mono_hessian_common::compute_single_op_hessian((*op).into(), x) + return common::compute_single_op_hessian((*op).into(), x) .expect("all single ops supported"); } diff --git a/src/mono/mono_ad_rr.rs b/src/mono/second_order/rr.rs similarity index 98% rename from src/mono/mono_ad_rr.rs rename to src/mono/second_order/rr.rs index 34dff7a..5091be0 100644 --- a/src/mono/mono_ad_rr.rs +++ b/src/mono/second_order/rr.rs @@ -142,8 +142,8 @@ use crate::Result; -use super::mono_hessian_common::{self, MonoHessianOpKind}; -use super::types::*; +use super::common::{self, MonoHessianOpKind}; +use crate::mono::types::*; /// Single-variable automatic differentiation operations for Reverse-over-Reverse Hessian computation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -175,7 +175,7 @@ impl MonoAD2RR { #[inline(always)] fn check_domain(self, x: f64) -> Result<()> { - mono_hessian_common::check_domain(self.as_hessian_op(), x) + common::check_domain(self.as_hessian_op(), x) } /// Forward pass computing value, first derivative, and second derivative. @@ -317,7 +317,7 @@ impl MonoAD2RR { // f'(x) = sign(x), with raw convention f'(0) = 0 // f''(x) = 0 away from zero; raw convention f''(0) = 0 let y = x.abs(); - let dy = mono_hessian_common::sign_or_zero(x); + let dy = common::sign_or_zero(x); let ddy = 0.0; (y, dy, ddy) } @@ -421,7 +421,7 @@ impl MonoAD2RR { Box::new(move |dy: f64| -> f64 { dy / (2.0 * sqrt_x) }) } MonoAD2RR::Abs => { - let sign = mono_hessian_common::sign_or_zero(x); + let sign = common::sign_or_zero(x); Box::new(move |dy: f64| -> f64 { dy * sign }) } }; diff --git a/src/mono/tests.rs b/src/mono/tests.rs index aeb621f..9353599 100644 --- a/src/mono/tests.rs +++ b/src/mono/tests.rs @@ -1,6 +1,6 @@ -use super::mono_ad::MonoAD; -use super::mono_fn::MonoFn; -use super::{mf1::MF1, mf2::MF2, mf3::MF3, mf4::MF4}; +use super::examples::{MF1, MF2, MF3, MF4}; +use super::first_order::MonoAD; +use super::func::MonoFn; use crate::mono_ops; use crate::test_utils::approx_eq_eps as approx_eq; use crate::AutodiffError; diff --git a/src/mono/tests_ho.rs b/src/mono/tests_ho.rs index 9f08908..aaefe6a 100644 --- a/src/mono/tests_ho.rs +++ b/src/mono/tests_ho.rs @@ -1,8 +1,8 @@ //! Tests for higher-order autodiff methods (RR, FR, RF). -use super::mono_ad_fr::MonoAD2FR; -use super::mono_ad_rf::MonoAD2RF; -use super::mono_ad_rr::MonoAD2RR; +use super::second_order::fr::MonoAD2FR; +use super::second_order::rf::MonoAD2RF; +use super::second_order::rr::MonoAD2RR; use crate::test_utils::approx_eq_eps as approx_eq; // Common tolerance for exact autodiff (machine precision) diff --git a/src/multi/backend/device.rs b/src/multi/compiled/backend/device.rs similarity index 99% rename from src/multi/backend/device.rs rename to src/multi/compiled/backend/device.rs index 08b400c..6a26915 100644 --- a/src/multi/backend/device.rs +++ b/src/multi/compiled/backend/device.rs @@ -1,11 +1,11 @@ //! Device-oriented buffer types, batch planning, and the DeviceBackend trait. -use crate::multi::backend::dispatch::BackendKind; +use crate::multi::compiled::backend::dispatch::BackendKind; use crate::multi::compiled::CompiledGraph; use crate::{AutodiffError, Result}; #[cfg(feature = "backend-wgpu")] -use crate::multi::backend::wgpu::WgpuBackend; +use crate::multi::compiled::backend::wgpu::WgpuBackend; /// Batch memory layout exposed for backend/device planning. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/multi/backend/dispatch.rs b/src/multi/compiled/backend/dispatch.rs similarity index 97% rename from src/multi/backend/dispatch.rs rename to src/multi/compiled/backend/dispatch.rs index c1d413e..19df276 100644 --- a/src/multi/backend/dispatch.rs +++ b/src/multi/compiled/backend/dispatch.rs @@ -1,12 +1,12 @@ //! Backend dispatch types, ExecutionBackend trait, and auto-dispatch logic. -use crate::multi::backend::device::{DeviceBackend, DeviceMemoryLocation}; -use crate::multi::backend::scalar::ScalarBackend; -use crate::multi::backend::simd::{ +use crate::multi::compiled::backend::device::{DeviceBackend, DeviceMemoryLocation}; +use crate::multi::compiled::backend::scalar::ScalarBackend; +use crate::multi::compiled::backend::simd::{ compute_batch_simd_f64x2, compute_batch_simd_f64x4, gradient_batch_simd_f64x2, gradient_batch_simd_f64x4, }; -use crate::multi::backend::types::{ +use crate::multi::compiled::backend::types::{ supports_simd_f64x2_runtime, supports_simd_f64x4_runtime, supports_wgpu_runtime, BackendCapabilities, OpCode, }; @@ -14,7 +14,7 @@ use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffe use crate::{AutodiffError, Result}; #[cfg(feature = "backend-wgpu")] -use crate::multi::backend::wgpu::WgpuBackend; +use crate::multi::compiled::backend::wgpu::WgpuBackend; /// Backend selected by automatic compiled-graph dispatch. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/multi/backend/mock.rs b/src/multi/compiled/backend/mock.rs similarity index 96% rename from src/multi/backend/mock.rs rename to src/multi/compiled/backend/mock.rs index 46d497f..0524ceb 100644 --- a/src/multi/backend/mock.rs +++ b/src/multi/compiled/backend/mock.rs @@ -1,12 +1,12 @@ //! Mock device-style backend that executes on CPU while using device-oriented plans. -use crate::multi::backend::device::{ +use crate::multi::compiled::backend::device::{ DeviceBackend, DeviceBufferKind, DeviceBufferSet, DeviceExecutionMode, DeviceExecutionTrace, DeviceTransferKind, }; -use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; -use crate::multi::backend::scalar::ScalarBackend; -use crate::multi::backend::types::BackendCapabilities; +use crate::multi::compiled::backend::dispatch::{BackendKind, ExecutionBackend}; +use crate::multi::compiled::backend::scalar::ScalarBackend; +use crate::multi::compiled::backend::types::BackendCapabilities; use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; use crate::{AutodiffError, Result}; diff --git a/src/multi/backend/mod.rs b/src/multi/compiled/backend/mod.rs similarity index 100% rename from src/multi/backend/mod.rs rename to src/multi/compiled/backend/mod.rs diff --git a/src/multi/backend/scalar.rs b/src/multi/compiled/backend/scalar.rs similarity index 91% rename from src/multi/backend/scalar.rs rename to src/multi/compiled/backend/scalar.rs index 032e245..8c3cbb7 100644 --- a/src/multi/backend/scalar.rs +++ b/src/multi/compiled/backend/scalar.rs @@ -1,8 +1,8 @@ //! Reference scalar f64 backend. -use crate::multi::backend::device::DeviceBackend; -use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; -use crate::multi::backend::types::BackendCapabilities; +use crate::multi::compiled::backend::device::DeviceBackend; +use crate::multi::compiled::backend::dispatch::{BackendKind, ExecutionBackend}; +use crate::multi::compiled::backend::types::BackendCapabilities; use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; use crate::{AutodiffError, Result}; diff --git a/src/multi/backend/simd.rs b/src/multi/compiled/backend/simd.rs similarity index 99% rename from src/multi/backend/simd.rs rename to src/multi/compiled/backend/simd.rs index f082427..98f695d 100644 --- a/src/multi/backend/simd.rs +++ b/src/multi/compiled/backend/simd.rs @@ -1,8 +1,8 @@ //! Prototype f64x2/f64x4 SIMD backend for batch compute and batch gradients. -use crate::multi::backend::device::DeviceBackend; -use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; -use crate::multi::backend::types::{ +use crate::multi::compiled::backend::device::DeviceBackend; +use crate::multi::compiled::backend::dispatch::{BackendKind, ExecutionBackend}; +use crate::multi::compiled::backend::types::{ supports_simd_f64x4_runtime, BackendCapabilities, FlatInstruction, OpCode, }; use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; diff --git a/src/multi/backend/types.rs b/src/multi/compiled/backend/types.rs similarity index 99% rename from src/multi/backend/types.rs rename to src/multi/compiled/backend/types.rs index 3032085..9e7512c 100644 --- a/src/multi/backend/types.rs +++ b/src/multi/compiled/backend/types.rs @@ -1,6 +1,6 @@ //! Core type definitions for the backend abstraction. -use crate::multi::multi_ad::MultiAD; +use crate::multi::first_order::MultiAD; use crate::{AutodiffError, NodeId, Result}; /// One closure-free instruction in a compiled scalar graph. diff --git a/src/multi/backend/wgpu.rs b/src/multi/compiled/backend/wgpu.rs similarity index 99% rename from src/multi/backend/wgpu.rs rename to src/multi/compiled/backend/wgpu.rs index 2b57109..541fc17 100644 --- a/src/multi/backend/wgpu.rs +++ b/src/multi/compiled/backend/wgpu.rs @@ -9,17 +9,17 @@ use pollster::block_on; use wgpu::{self, BufferUsages}; #[cfg(feature = "backend-wgpu")] -use crate::multi::backend::device::{ +use crate::multi::compiled::backend::device::{ AcceleratorDeviceContext, AcceleratorDeviceKind, DeviceBackend, DeviceBatchPlan, DeviceBufferHandle, DeviceBufferKind, DeviceExecutionMode, DeviceExecutionTrace, DeviceTransferKind, DeviceTransferPolicy, GpuBackendBoundary, }; #[cfg(feature = "backend-wgpu")] -use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; +use crate::multi::compiled::backend::dispatch::{BackendKind, ExecutionBackend}; #[cfg(feature = "backend-wgpu")] -use crate::multi::backend::scalar::ScalarBackend; +use crate::multi::compiled::backend::scalar::ScalarBackend; #[cfg(feature = "backend-wgpu")] -use crate::multi::backend::types::{BackendCapabilities, OpCode, UNUSED_NODE_ID}; +use crate::multi::compiled::backend::types::{BackendCapabilities, OpCode, UNUSED_NODE_ID}; #[cfg(feature = "backend-wgpu")] use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; #[cfg(feature = "backend-wgpu")] diff --git a/src/multi/compiled/batches.rs b/src/multi/compiled/batches.rs new file mode 100644 index 0000000..9e1e739 --- /dev/null +++ b/src/multi/compiled/batches.rs @@ -0,0 +1,221 @@ +//! Batch data types for evaluating compiled graphs on multiple inputs. + +use crate::{AutodiffError, Result}; + +/// Flat row-major batch input view. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct BatchInputs<'a> { + /// Flat row-major input data. + pub data: &'a [f64], + /// Number of rows in the batch. + pub batch_size: usize, + /// Number of inputs per row. + pub input_dim: usize, +} + +impl<'a> BatchInputs<'a> { + /// Create a validated row-major input view. + pub fn new(data: &'a [f64], batch_size: usize, input_dim: usize) -> Result { + if data.len() != batch_size.saturating_mul(input_dim) { + return Err(AutodiffError::InvalidGraph { + reason: "batch data length must equal batch_size * input_dim", + }); + } + Ok(Self { + data, + batch_size, + input_dim, + }) + } + + /// Return one input row. + /// + /// Prefer [`BatchInputs::try_row`] when invalid row indices should return an error. + #[must_use] + pub fn row(&self, index: usize) -> &[f64] { + self.try_row(index).expect("batch row index out of bounds") + } + + /// Return one input row, or an error when `index` is out of range. + pub fn try_row(&self, index: usize) -> Result<&[f64]> { + if index >= self.batch_size { + return Err(AutodiffError::IndexOutOfBounds { + index, + max_index: self.batch_size.saturating_sub(1), + }); + } + let start = index + .checked_mul(self.input_dim) + .ok_or(AutodiffError::IndexOutOfBounds { + index, + max_index: self.batch_size.saturating_sub(1), + })?; + let end = start + .checked_add(self.input_dim) + .ok_or(AutodiffError::IndexOutOfBounds { + index, + max_index: self.batch_size.saturating_sub(1), + })?; + self.data + .get(start..end) + .ok_or(AutodiffError::IndexOutOfBounds { + index, + max_index: self.batch_size.saturating_sub(1), + }) + } +} + +/// Flat row-major batch values. +#[derive(Debug, Clone, PartialEq)] +pub struct BatchValues { + pub data: Vec, + pub batch_size: usize, + pub output_dim: usize, +} + +/// Reusable flat row-major batch value buffer. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BatchValuesBuffer { + pub data: Vec, + pub batch_size: usize, + pub output_dim: usize, +} + +impl BatchValuesBuffer { + /// Create an empty reusable output buffer. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + pub(crate) fn reset(&mut self, batch_size: usize, output_dim: usize) { + self.data.clear(); + self.data.reserve(batch_size.saturating_mul(output_dim)); + self.batch_size = batch_size; + self.output_dim = output_dim; + } + + /// Clone the current buffer contents into an owned result value. + #[must_use] + pub fn to_values(&self) -> BatchValues { + BatchValues { + data: self.data.clone(), + batch_size: self.batch_size, + output_dim: self.output_dim, + } + } +} + +/// Flat row-major batch scalar-output gradients. +#[derive(Debug, Clone, PartialEq)] +pub struct BatchGradients { + pub values: Vec, + pub gradients: Vec, + pub batch_size: usize, + pub input_dim: usize, +} + +/// Reusable flat row-major batch scalar-output gradient buffer. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BatchGradientsBuffer { + pub values: Vec, + pub gradients: Vec, + pub batch_size: usize, + pub input_dim: usize, +} + +impl BatchGradientsBuffer { + /// Create an empty reusable gradient buffer. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + pub(crate) fn reset(&mut self, batch_size: usize, input_dim: usize) { + self.values.clear(); + self.gradients.clear(); + self.values.reserve(batch_size); + self.gradients.reserve(batch_size.saturating_mul(input_dim)); + self.batch_size = batch_size; + self.input_dim = input_dim; + } + + /// Clone the current buffer contents into an owned result value. + #[must_use] + pub fn to_gradients(&self) -> BatchGradients { + BatchGradients { + values: self.values.clone(), + gradients: self.gradients.clone(), + batch_size: self.batch_size, + input_dim: self.input_dim, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn test_batch_inputs_try_row_valid() { + let data = vec![1.0, 2.0, 3.0, 4.0]; + let batch = BatchInputs::new(&data, 2, 2).unwrap(); + let row0 = batch.try_row(0).unwrap(); + assert_eq!(row0, &[1.0, 2.0]); + let row1 = batch.try_row(1).unwrap(); + assert_eq!(row1, &[3.0, 4.0]); + } + + #[test] + fn test_batch_inputs_try_row_out_of_bounds() { + let data = vec![1.0, 2.0, 3.0, 4.0]; + let batch = BatchInputs::new(&data, 2, 2).unwrap(); + assert!(batch.try_row(2).is_err()); + assert!(batch.try_row(100).is_err()); + } + + #[test] + fn test_batch_inputs_invalid_shape() { + assert!(BatchInputs::new(&[1.0, 2.0], 2, 2).is_err()); + assert!(BatchInputs::new(&[1.0, 2.0, 3.0], 1, 2).is_err()); + } + + #[test] + fn test_batch_values_buffer_reuse() { + // We can't test with CompiledGraph here — just test buffer lifecycle. + let mut buffer = BatchValuesBuffer::new(); + + buffer.reset(1, 1); + buffer.data.push(3.0); + let values = buffer.to_values(); + assert_eq!(values.batch_size, 1); + assert_eq!(values.data, &[3.0]); + + buffer.reset(2, 2); + buffer.data.extend_from_slice(&[9.0, 13.0]); + let values2 = buffer.to_values(); + assert_eq!(values2.batch_size, 2); + assert_eq!(values2.data, &[9.0, 13.0]); + } + + #[test] + fn test_batch_gradients_buffer_reuse() { + let mut buffer = BatchGradientsBuffer::new(); + + buffer.reset(1, 2); + buffer.values.push(3.0); + buffer.gradients.extend_from_slice(&[1.0, 1.0]); + let grad = buffer.to_gradients(); + assert_eq!(grad.batch_size, 1); + assert_eq!(grad.values, &[3.0]); + assert_eq!(grad.gradients, &[1.0, 1.0]); + + buffer.reset(2, 2); + buffer.values.push(9.0); + buffer.values.push(13.0); + buffer.gradients.extend_from_slice(&[1.0, 1.0, 1.0, 1.0]); + let grad2 = buffer.to_gradients(); + assert_eq!(grad2.batch_size, 2); + assert_eq!(grad2.values, &[9.0, 13.0]); + assert_eq!(grad2.gradients, &[1.0, 1.0, 1.0, 1.0]); + } +} diff --git a/src/multi/compiled.rs b/src/multi/compiled/ir.rs similarity index 91% rename from src/multi/compiled.rs rename to src/multi/compiled/ir.rs index 1e683ab..21f04b2 100644 --- a/src/multi/compiled.rs +++ b/src/multi/compiled/ir.rs @@ -1,8 +1,5 @@ //! Closure-free compiled instruction IR for acceleration-ready graph execution. -use super::op_rules; -use crate::{AutodiffError, NodeId, Result}; - use super::backend::{ BackendCapabilities, BackendKind, BackendRejectionReason, BackendSupportReport, DeviceBackend, DeviceBatchPlan, DeviceBufferSet, DeviceExecutionTrace, FlatInstruction, Instruction, @@ -12,142 +9,13 @@ use super::backend::{ use super::backend::{DeviceBufferKind, DeviceMemoryLocation}; #[cfg(feature = "backend-wgpu")] use super::backend::{WgpuBackend, WgpuBufferSet}; +use super::batches::{ + BatchGradients, BatchGradientsBuffer, BatchInputs, BatchValues, BatchValuesBuffer, +}; #[cfg(test)] -use super::multi_ad::MultiAD; - -/// Flat row-major batch input view. -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct BatchInputs<'a> { - /// Flat row-major input data. - pub data: &'a [f64], - /// Number of rows in the batch. - pub batch_size: usize, - /// Number of inputs per row. - pub input_dim: usize, -} - -impl<'a> BatchInputs<'a> { - /// Create a validated row-major input view. - pub fn new(data: &'a [f64], batch_size: usize, input_dim: usize) -> Result { - if data.len() != batch_size.saturating_mul(input_dim) { - return Err(AutodiffError::InvalidGraph { - reason: "batch data length must equal batch_size * input_dim", - }); - } - Ok(Self { - data, - batch_size, - input_dim, - }) - } - - /// Return one input row. - /// - /// Prefer [`BatchInputs::try_row`] when invalid row indices should return an error. - #[must_use] - pub fn row(&self, index: usize) -> &[f64] { - self.try_row(index).expect("batch row index out of bounds") - } - - /// Return one input row, or an error when `index` is out of range. - pub fn try_row(&self, index: usize) -> Result<&[f64]> { - if index >= self.batch_size { - return Err(AutodiffError::IndexOutOfBounds { - index, - max_index: self.batch_size.saturating_sub(1), - }); - } - let start = index * self.input_dim; - Ok(&self.data[start..start + self.input_dim]) - } -} - -/// Flat row-major batch values. -#[derive(Debug, Clone, PartialEq)] -pub struct BatchValues { - pub data: Vec, - pub batch_size: usize, - pub output_dim: usize, -} - -/// Reusable flat row-major batch value buffer. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct BatchValuesBuffer { - pub data: Vec, - pub batch_size: usize, - pub output_dim: usize, -} - -impl BatchValuesBuffer { - /// Create an empty reusable output buffer. - #[must_use] - pub fn new() -> Self { - Self::default() - } - - pub(crate) fn reset(&mut self, batch_size: usize, output_dim: usize) { - self.data.clear(); - self.data.reserve(batch_size.saturating_mul(output_dim)); - self.batch_size = batch_size; - self.output_dim = output_dim; - } - - /// Clone the current buffer contents into an owned result value. - #[must_use] - pub fn to_values(&self) -> BatchValues { - BatchValues { - data: self.data.clone(), - batch_size: self.batch_size, - output_dim: self.output_dim, - } - } -} - -/// Flat row-major batch scalar-output gradients. -#[derive(Debug, Clone, PartialEq)] -pub struct BatchGradients { - pub values: Vec, - pub gradients: Vec, - pub batch_size: usize, - pub input_dim: usize, -} - -/// Reusable flat row-major batch scalar-output gradient buffer. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct BatchGradientsBuffer { - pub values: Vec, - pub gradients: Vec, - pub batch_size: usize, - pub input_dim: usize, -} - -impl BatchGradientsBuffer { - /// Create an empty reusable gradient buffer. - #[must_use] - pub fn new() -> Self { - Self::default() - } - - pub(crate) fn reset(&mut self, batch_size: usize, input_dim: usize) { - self.values.clear(); - self.gradients.clear(); - self.values.reserve(batch_size); - self.gradients.reserve(batch_size.saturating_mul(input_dim)); - self.batch_size = batch_size; - self.input_dim = input_dim; - } - - /// Clone the current buffer contents into an owned result value. - #[must_use] - pub fn to_gradients(&self) -> BatchGradients { - BatchGradients { - values: self.values.clone(), - gradients: self.gradients.clone(), - batch_size: self.batch_size, - input_dim: self.input_dim, - } - } -} +use crate::multi::first_order::MultiAD; +use crate::multi::op_rules; +use crate::{AutodiffError, NodeId, Result}; /// Static metadata for backend selection and workspace planning. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/multi/compiled/mod.rs b/src/multi/compiled/mod.rs new file mode 100644 index 0000000..7447027 --- /dev/null +++ b/src/multi/compiled/mod.rs @@ -0,0 +1,10 @@ +//! Closure-free compiled instruction IR for acceleration-ready graph execution. + +pub mod backend; +pub mod batches; +pub mod ir; + +pub use batches::{ + BatchGradients, BatchGradientsBuffer, BatchInputs, BatchValues, BatchValuesBuffer, +}; +pub use ir::{CompiledGraph, CompiledGraphMetadata, CompiledWorkspace}; diff --git a/src/multi/examples.rs b/src/multi/examples.rs new file mode 100644 index 0000000..2625493 --- /dev/null +++ b/src/multi/examples.rs @@ -0,0 +1,116 @@ +//! Test-only example implementations of [`MultiFn`] for library tests. +//! +//! These are production-type examples used internally by the test suite and +//! are gated behind `#[cfg(test)]` to keep them out of release builds. + +#[cfg(test)] +use super::first_order::MultiAD; +#[cfg(test)] +use super::func::{GraphType, MultiFn}; +#[cfg(test)] +use crate::multi_ops; + +/// f(x₁, x₂) = sin(x₁) * (x₁ + x₂) +#[cfg(test)] +pub struct F1(pub f64, pub f64); + +#[cfg(test)] +impl MultiFn for F1 { + fn inputs(&self) -> Vec { + vec![self.0, self.1] + } + + fn graph(&self) -> &'static GraphType { + use std::sync::LazyLock; + static GRAPH: LazyLock)>> = LazyLock::new(|| { + Vec::from(multi_ops![ + (inp, 0), // x₁ at index 0 + (inp, 1), // x₂ at index 1 + (add, 0, 1), // x₁ + x₂ at index 2 + (sin, 0), // sin(x₁) at index 3 + (mul, 2, 3), // sin(x₁) * (x₁ + x₂) at index 4 + ]) + }); + &GRAPH + } + + fn expected_value(&self) -> f64 { + self.0.sin() * (self.0 + self.1) + } + + fn expected_gradients(&self) -> Vec { + let df_dx1 = self.0.cos() * (self.0 + self.1) + self.0.sin(); + let df_dx2 = self.0.sin(); + vec![df_dx1, df_dx2] + } +} + +/// f(x₁, x₂) = sin(x₁) / (x₁ - x₂) +#[cfg(test)] +pub struct F2(pub f64, pub f64); + +#[cfg(test)] +impl MultiFn for F2 { + fn inputs(&self) -> Vec { + vec![self.0, self.1] + } + + fn graph(&self) -> &'static GraphType { + use std::sync::LazyLock; + static GRAPH: LazyLock)>> = LazyLock::new(|| { + Vec::from(multi_ops![ + (inp, 0), // x₁ at index 0 + (inp, 1), // x₂ at index 1 + (sub, 0, 1), // x₁ - x₂ at index 2 + (sin, 0), // sin(x₁) at index 3 + (div, 3, 2), // sin(x₁) / (x₁ - x₂) at index 4 + ]) + }); + &GRAPH + } + + fn expected_value(&self) -> f64 { + self.0.sin() / (self.0 - self.1) + } + + fn expected_gradients(&self) -> Vec { + let df_dx1 = self.0.cos() / (self.0 - self.1) - self.0.sin() / (self.0 - self.1).powi(2); + let df_dx2 = self.0.sin() / (self.0 - self.1).powi(2); + vec![df_dx1, df_dx2] + } +} + +/// f(x₁, x₂) = sin(x₁) * ln(x₂) +#[cfg(test)] +pub struct F3(pub f64, pub f64); + +#[cfg(test)] +impl MultiFn for F3 { + fn inputs(&self) -> Vec { + vec![self.0, self.1] + } + + fn graph(&self) -> &'static GraphType { + use std::sync::LazyLock; + static GRAPH: LazyLock)>> = LazyLock::new(|| { + Vec::from(multi_ops![ + (inp, 0), // x₁ at index 0 + (inp, 1), // x₂ at index 1 + (ln, 1), // ln(x₂) at index 2 + (sin, 0), // sin(x₁) at index 3 + (mul, 3, 2), // sin(x₁) * ln(x₂) at index 4 + ]) + }); + &GRAPH + } + + fn expected_value(&self) -> f64 { + self.0.sin() * self.1.ln() + } + + fn expected_gradients(&self) -> Vec { + let df_dx1 = self.0.cos() * self.1.ln(); + let df_dx2 = self.0.sin() / self.1; + vec![df_dx1, df_dx2] + } +} diff --git a/src/multi/f1.rs b/src/multi/f1.rs deleted file mode 100644 index 2597fda..0000000 --- a/src/multi/f1.rs +++ /dev/null @@ -1,43 +0,0 @@ -#[cfg(test)] -use super::multi_ad::MultiAD; -#[cfg(test)] -use super::multi_fn::{GraphType, MultiFn}; -#[cfg(test)] -use crate::multi_ops; - -#[cfg(test)] -pub struct F1(pub f64, pub f64); // Represents f(x₁, x₂) = sin(x₁) * (x₁ + x₂) - -#[cfg(test)] -impl MultiFn for F1 { - fn inputs(&self) -> Vec { - vec![self.0, self.1] - } - - fn graph(&self) -> &'static GraphType { - use std::sync::LazyLock; - static GRAPH: LazyLock)>> = LazyLock::new(|| { - Vec::from(multi_ops![ - (inp, 0), // x₁ at index 0 - (inp, 1), // x₂ at index 1 - (add, 0, 1), // x₁ + x₂ at index 2 - (sin, 0), // sin(x₁) at index 3 - (mul, 2, 3), // sin(x₁) * (x₁ + x₂) at index 4 - ]) - }); - &GRAPH - } - - /// Example function: f(x₁, x₂) = sin(x₁) * (x₁ + x₂) - fn expected_value(&self) -> f64 { - self.0.sin() * (self.0 + self.1) - } - - /// Analytical gradient of f: (∂f/∂x₁, ∂f/∂x₂) - /// Using product rule: d(sin(x) * (x + y))/dx = cos(x) * (x + y) + sin(x) - fn expected_gradients(&self) -> Vec { - let df_dx1 = self.0.cos() * (self.0 + self.1) + self.0.sin(); - let df_dx2 = self.0.sin(); - vec![df_dx1, df_dx2] - } -} diff --git a/src/multi/f2.rs b/src/multi/f2.rs deleted file mode 100644 index fa377d6..0000000 --- a/src/multi/f2.rs +++ /dev/null @@ -1,43 +0,0 @@ -#[cfg(test)] -use super::multi_ad::MultiAD; -#[cfg(test)] -use super::multi_fn::{GraphType, MultiFn}; -#[cfg(test)] -use crate::multi_ops; - -#[cfg(test)] -pub struct F2(pub f64, pub f64); // Represents f(x₁, x₂) = sin(x₁) / (x₁ - x₂) - -#[cfg(test)] -impl MultiFn for F2 { - fn inputs(&self) -> Vec { - vec![self.0, self.1] - } - - fn graph(&self) -> &'static GraphType { - use std::sync::LazyLock; - static GRAPH: LazyLock)>> = LazyLock::new(|| { - Vec::from(multi_ops![ - (inp, 0), // x₁ at index 0 - (inp, 1), // x₂ at index 1 - (sub, 0, 1), // x₁ - x₂ at index 2 - (sin, 0), // sin(x₁) at index 3 - (div, 3, 2), // sin(x₁) / (x₁ - x₂) at index 4 - ]) - }); - &GRAPH - } - - /// Example function: f(x₁, x₂) = sin(x₁) / (x₁ - x₂) - fn expected_value(&self) -> f64 { - self.0.sin() / (self.0 - self.1) - } - - /// Analytical gradient of f: (∂f/∂x₁, ∂f/∂x₂) - /// Using quotient rule: d(sin(x) / (x - y))/dx = cos(x) / (x - y) - sin(x) / (x - y)² - fn expected_gradients(&self) -> Vec { - let df_dx1 = self.0.cos() / (self.0 - self.1) - self.0.sin() / (self.0 - self.1).powi(2); - let df_dx2 = self.0.sin() / (self.0 - self.1).powi(2); - vec![df_dx1, df_dx2] - } -} diff --git a/src/multi/f3.rs b/src/multi/f3.rs deleted file mode 100644 index 61bea23..0000000 --- a/src/multi/f3.rs +++ /dev/null @@ -1,43 +0,0 @@ -#[cfg(test)] -use super::multi_ad::MultiAD; -#[cfg(test)] -use super::multi_fn::{GraphType, MultiFn}; -#[cfg(test)] -use crate::multi_ops; - -#[cfg(test)] -pub struct F3(pub f64, pub f64); // Represents f(x₁, x₂) = sin(x₁) * ln(x₂) - -#[cfg(test)] -impl MultiFn for F3 { - fn inputs(&self) -> Vec { - vec![self.0, self.1] - } - - fn graph(&self) -> &'static GraphType { - use std::sync::LazyLock; - static GRAPH: LazyLock)>> = LazyLock::new(|| { - Vec::from(multi_ops![ - (inp, 0), // x₁ at index 0 - (inp, 1), // x₂ at index 1 - (ln, 1), // ln(x₂) at index 2 - (sin, 0), // sin(x₁) at index 3 - (mul, 3, 2), // sin(x₁) * ln(x₂) at index 4 - ]) - }); - &GRAPH - } - - /// Example function: f(x₁, x₂) = sin(x₁) * ln(x₂) - fn expected_value(&self) -> f64 { - self.0.sin() * self.1.ln() - } - - /// Analytical gradient of f: (∂f/∂x₁, ∂f/∂x₂) - /// Using product rule: d(sin(x) * ln(y))/dx = cos(x) * ln(y), d/dy = sin(x) / y - fn expected_gradients(&self) -> Vec { - let df_dx1 = self.0.cos() * self.1.ln(); - let df_dx2 = self.0.sin() / self.1; - vec![df_dx1, df_dx2] - } -} diff --git a/src/multi/multi_ad.rs b/src/multi/first_order.rs similarity index 100% rename from src/multi/multi_ad.rs rename to src/multi/first_order.rs diff --git a/src/multi/multi_fn.rs b/src/multi/func.rs similarity index 99% rename from src/multi/multi_fn.rs rename to src/multi/func.rs index 6fc91f8..bfcb566 100644 --- a/src/multi/multi_fn.rs +++ b/src/multi/func.rs @@ -1,4 +1,4 @@ -pub use super::multi_ad::MultiAD; +pub use super::first_order::MultiAD; pub use super::types::BackwardResultBox; use crate::error::Result; diff --git a/src/multi/builder.rs b/src/multi/graph/builder.rs similarity index 99% rename from src/multi/builder.rs rename to src/multi/graph/builder.rs index 86e67ab..31247fc 100644 --- a/src/multi/builder.rs +++ b/src/multi/graph/builder.rs @@ -3,8 +3,8 @@ //! This module provides a fluent, type-safe interface for building computational //! graphs without manually managing indices and vectors. -use super::graph::{Graph, GraphNode, NodeId}; -use super::multi_ad::MultiAD; +use crate::multi::first_order::MultiAD; +use crate::multi::graph::core::{Graph, GraphNode, NodeId}; /// Builder for constructing multi-variable computation graphs. /// @@ -445,7 +445,7 @@ impl GraphBuilder { #[cfg(test)] mod tests { use super::*; - use crate::multi::multi_ad::MultiAD; + use crate::multi::first_order::MultiAD; use crate::test_utils::approx_eq_eps as approx_eq; #[test] diff --git a/src/multi/graph.rs b/src/multi/graph/core.rs similarity index 99% rename from src/multi/graph.rs rename to src/multi/graph/core.rs index bd69aef..0ccaf17 100644 --- a/src/multi/graph.rs +++ b/src/multi/graph/core.rs @@ -35,25 +35,25 @@ use std::{fmt::Write, sync::Arc}; -use super::backend::{ +use crate::multi::compiled::backend::{ BackendKind, BackendSupportReport, DeviceBatchPlan, DeviceBufferSet, DeviceExecutionTrace, Instruction, }; -use super::compiled::{ +use crate::multi::compiled::{ BatchGradients, BatchGradientsBuffer, BatchInputs, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, CompiledWorkspace, }; +use crate::multi::first_order::MultiAD; #[cfg(test)] -use super::expr::ExprGraph; -use super::multi_ad::MultiAD; -use super::multi_ad_fr::MultiAD2FR; -use super::multi_ad_rf::MultiAD2RF; -use super::op_rules; -use super::parser; +use crate::multi::graph::expr::ExprGraph; +use crate::multi::graph::parser; #[cfg(test)] -use super::tape::TapeWorkspace; -use super::tape::{CompiledArgRange, CompiledNode, Tape}; -use super::types::BackwardResultBox; +use crate::multi::graph::tape::TapeWorkspace; +use crate::multi::graph::tape::{CompiledArgRange, CompiledNode, Tape}; +use crate::multi::op_rules; +use crate::multi::second_order::fr::MultiAD2FR; +use crate::multi::second_order::rf::MultiAD2RF; +use crate::multi::types::BackwardResultBox; use crate::{AutodiffError, Result}; /// A handle to an input or computed node in a graph. diff --git a/src/multi/expr.rs b/src/multi/graph/expr.rs similarity index 98% rename from src/multi/expr.rs rename to src/multi/graph/expr.rs index 2737020..cc94fcc 100644 --- a/src/multi/expr.rs +++ b/src/multi/graph/expr.rs @@ -6,8 +6,8 @@ use std::{ rc::Rc, }; -use super::graph::{Graph, NodeId}; -use super::multi_ad::MultiAD; +use crate::multi::first_order::MultiAD; +use crate::multi::graph::core::{Graph, NodeId}; use crate::Result; /// Shared expression graph used for operator-overloaded graph construction. diff --git a/src/multi/graph/mod.rs b/src/multi/graph/mod.rs new file mode 100644 index 0000000..31fe521 --- /dev/null +++ b/src/multi/graph/mod.rs @@ -0,0 +1,7 @@ +//! Reusable graph API and tape-based evaluation for multi-variable autodiff. + +pub mod builder; +pub mod core; +pub mod expr; +pub mod parser; +pub mod tape; diff --git a/src/multi/parser.rs b/src/multi/graph/parser.rs similarity index 100% rename from src/multi/parser.rs rename to src/multi/graph/parser.rs diff --git a/src/multi/tape.rs b/src/multi/graph/tape.rs similarity index 99% rename from src/multi/tape.rs rename to src/multi/graph/tape.rs index 134aef0..e7dae62 100644 --- a/src/multi/tape.rs +++ b/src/multi/graph/tape.rs @@ -2,10 +2,10 @@ use std::sync::Arc; -use super::graph::{Graph, NodeId}; -use super::multi_ad::MultiAD; -use super::op_rules; -use super::types::BackwardResultBox; +use crate::multi::first_order::MultiAD; +use crate::multi::graph::core::{Graph, NodeId}; +use crate::multi::op_rules; +use crate::multi::types::BackwardResultBox; use crate::{AutodiffError, Result}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/multi.rs b/src/multi/mod.rs similarity index 61% rename from src/multi.rs rename to src/multi/mod.rs index 29026b1..d45a218 100644 --- a/src/multi.rs +++ b/src/multi/mod.rs @@ -3,32 +3,27 @@ //! This module provides functionality for computing gradients of //! multi-variable functions using computational graphs. -mod f1; -mod f2; -mod f3; - -pub mod backend; -pub mod builder; pub mod compiled; -pub mod expr; +pub mod first_order; +pub mod func; pub mod graph; -mod multi_ad; -pub mod multi_ad_fr; -pub mod multi_ad_rf; -pub mod multi_ad_rr; -mod parser; -pub mod tape; - -// Shared internal modules for multivariate derivative rules and Hessian computation. -mod multi_hessian_common; pub(crate) mod op_rules; +pub mod second_order; +pub mod types; -mod multi_fn; +#[cfg(test)] +mod examples; #[cfg(test)] mod tests; -pub mod types; -pub use backend::{ +pub use first_order::MultiAD; +pub use func::MultiFn; + +pub use second_order::fr::MultiAD2FR; +pub use second_order::rf::MultiAD2RF; +pub use second_order::rr::MultiAD2RR; + +pub use compiled::backend::{ AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, BackendKind, BackendRejectionReason, BackendSupportReport, BatchLayout, DeviceBackend, DeviceBatchPlan, DeviceBuffer, DeviceBufferHandle, DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, @@ -38,24 +33,16 @@ pub use backend::{ UNUSED_NODE_ID, }; #[cfg(feature = "backend-wgpu")] -pub use backend::{ +pub use compiled::backend::{ WgpuBackend, WgpuBuffer, WgpuBufferSet, WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES, }; pub use compiled::{ BatchGradients, BatchGradientsBuffer, BatchInputs, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, CompiledWorkspace, }; -pub use expr::{ExprGraph, ExprNode}; -pub use graph::{ +pub use graph::builder::GraphBuilder; +pub use graph::core::{ DomainPolicy, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, GraphStats, NodeId, }; -pub use multi_ad::MultiAD; -pub use multi_fn::MultiFn; -pub use tape::{Tape, TapeWorkspace}; - -#[allow(unused_imports)] // Used via lib.rs re-export path -pub use multi_ad_fr::MultiAD2FR; -#[allow(unused_imports)] // Used via lib.rs re-export path -pub use multi_ad_rf::MultiAD2RF; -#[allow(unused_imports)] // Used via lib.rs re-export path -pub use multi_ad_rr::MultiAD2RR; +pub use graph::expr::{ExprGraph, ExprNode}; +pub use graph::tape::{Tape, TapeWorkspace}; diff --git a/src/multi/op_rules.rs b/src/multi/op_rules.rs index 27ece17..a205781 100644 --- a/src/multi/op_rules.rs +++ b/src/multi/op_rules.rs @@ -1,10 +1,10 @@ //! Shared local operation rules for multivariate autodiff. //! //! This module centralizes scalar values plus first- and second-order local -//! derivatives for [`super::multi_ad::MultiAD`] operations so first-order, +//! derivatives for [`super::first_order::MultiAD`] operations so first-order, //! forward-mode, and exact Hessian implementations can reuse the same formulas. -use super::multi_ad::MultiAD; +use super::first_order::MultiAD; use crate::{AutodiffError, Result}; #[inline(always)] diff --git a/src/multi/multi_hessian_common.rs b/src/multi/second_order/common.rs similarity index 99% rename from src/multi/multi_hessian_common.rs rename to src/multi/second_order/common.rs index 88639ab..1f7b8bb 100644 --- a/src/multi/multi_hessian_common.rs +++ b/src/multi/second_order/common.rs @@ -14,8 +14,8 @@ use crate::error::{AutodiffError, Result}; -use super::multi_ad::MultiAD; -use super::op_rules::{self, LocalRule}; +use crate::multi::first_order::MultiAD; +use crate::multi::op_rules::{self, LocalRule}; /// An operation descriptor for the stack-based RPN computation graph. /// Used by FR and RF Hessian methods. diff --git a/src/multi/multi_ad_fr.rs b/src/multi/second_order/fr.rs similarity index 99% rename from src/multi/multi_ad_fr.rs rename to src/multi/second_order/fr.rs index 4bafc85..ba4b4bb 100644 --- a/src/multi/multi_ad_fr.rs +++ b/src/multi/second_order/fr.rs @@ -44,7 +44,7 @@ use std::fmt; use crate::Result; -use super::multi_hessian_common::{compute_hessian_dual, OpKind}; +use super::common::{compute_hessian_dual, OpKind}; /// Stack-based operation for multivariate second-order AD. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] diff --git a/src/multi/second_order/mod.rs b/src/multi/second_order/mod.rs new file mode 100644 index 0000000..6f6bd6c --- /dev/null +++ b/src/multi/second_order/mod.rs @@ -0,0 +1,6 @@ +//! Exact second-order (Hessian) computation for multi-variable functions. + +pub(crate) mod common; +pub mod fr; +pub mod rf; +pub mod rr; diff --git a/src/multi/multi_ad_rf.rs b/src/multi/second_order/rf.rs similarity index 99% rename from src/multi/multi_ad_rf.rs rename to src/multi/second_order/rf.rs index 66aa171..3015b11 100644 --- a/src/multi/multi_ad_rf.rs +++ b/src/multi/second_order/rf.rs @@ -56,7 +56,7 @@ use std::fmt; use crate::Result; -use super::multi_hessian_common::{compute_hessian_dual, OpKind}; +use super::common::{compute_hessian_dual, OpKind}; /// Stack-based operation for multivariate second-order AD. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] diff --git a/src/multi/multi_ad_rr.rs b/src/multi/second_order/rr.rs similarity index 99% rename from src/multi/multi_ad_rr.rs rename to src/multi/second_order/rr.rs index 49f5479..059abe2 100644 --- a/src/multi/multi_ad_rr.rs +++ b/src/multi/second_order/rr.rs @@ -36,8 +36,8 @@ use std::fmt; use crate::error::{AutodiffError, Result}; -use super::multi_ad::MultiAD; -use super::op_rules::{self, LocalRule}; +use crate::multi::first_order::MultiAD; +use crate::multi::op_rules::{self, LocalRule}; /// Stack-based operation for multivariate second-order AD. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] diff --git a/src/multi/tests.rs b/src/multi/tests.rs index c53813d..5fc40d6 100644 --- a/src/multi/tests.rs +++ b/src/multi/tests.rs @@ -1,6 +1,6 @@ -use super::f1::F1; -use super::f2::F2; -use super::f3::F3; +use super::examples::F1; +use super::examples::F2; +use super::examples::F3; use super::*; use crate::multi_ops; use crate::test_utils::approx_eq_eps as approx_eq;