pub trait SymmetricQuadraticCoefficients {
fn dimension(&self) -> usize;
fn multiply(&self, input: &[f64], output: &mut [f64]);
fn coefficient(&self, row: usize, column: usize) -> f64;
fn visit_upper_triangle(
&self,
direction: &mut [f64],
projected: &mut [f64],
mut visit: impl FnMut(usize, usize, f64),
) {
let dimension = self.dimension();
assert_eq!(direction.len(), dimension);
assert_eq!(projected.len(), dimension);
direction.fill(0.0);
for column in 0..dimension {
direction[column] = 1.0;
self.multiply(direction, projected);
direction[column] = 0.0;
for row in 0..=column {
visit(row, column, projected[row]);
}
}
}
fn quadratic_value<T, F>(&self, inputs: &[T], value: F) -> f64
where
F: Fn(&T) -> f64,
{
assert_eq!(
inputs.len(),
self.dimension(),
"symmetric quadratic-form dimension mismatch"
);
let mut out = 0.0;
for row in 0..inputs.len() {
let row_value = value(&inputs[row]);
out += self.coefficient(row, row) * row_value * row_value;
for column in row + 1..inputs.len() {
out += 2.0 * self.coefficient(row, column) * row_value * value(&inputs[column]);
}
}
out
}
}
fn symmetric_quadratic_form_default<T, C>(
inputs: &[T],
coefficients: &C,
constant: impl Fn(f64) -> T,
add: impl Fn(&T, &T) -> T,
mul: impl Fn(&T, &T) -> T,
scale: impl Fn(&T, f64) -> T,
) -> T
where
C: SymmetricQuadraticCoefficients,
{
assert_eq!(
inputs.len(),
coefficients.dimension(),
"symmetric quadratic-form dimension mismatch"
);
let mut out = constant(0.0);
for row in 0..inputs.len() {
let diagonal = mul(&inputs[row], &inputs[row]);
out = add(&out, &scale(&diagonal, coefficients.coefficient(row, row)));
for column in row + 1..inputs.len() {
let cross = mul(&inputs[row], &inputs[column]);
out = add(
&out,
&scale(&cross, 2.0 * coefficients.coefficient(row, column)),
);
}
}
out
}
fn linear_combination_default<T>(
inputs: &[T],
weights: &[f64],
constant: impl Fn(f64) -> T,
add: impl Fn(&T, &T) -> T,
scale: impl Fn(&T, f64) -> T,
) -> T {
assert_eq!(
inputs.len(),
weights.len(),
"linear-combination dimension mismatch"
);
inputs
.iter()
.zip(weights)
.fold(constant(0.0), |sum, (input, &weight)| {
add(&sum, &scale(input, weight))
})
}
fn multiply_add_default<T>(
left: &T,
right: &T,
addend: &T,
mul: impl Fn(&T, &T) -> T,
add: impl Fn(&T, &T) -> T,
) -> T {
add(&mul(left, right), addend)
}
fn composed_sum_default<T>(
inputs: &[T],
derivative_stacks: &[[f64; 5]],
constant: impl Fn(f64) -> T,
add: impl Fn(&T, &T) -> T,
compose: impl Fn(&T, [f64; 5]) -> T,
) -> T {
assert_eq!(
inputs.len(),
derivative_stacks.len(),
"composed-sum term-count mismatch"
);
inputs
.iter()
.zip(derivative_stacks)
.fold(constant(0.0), |sum, (input, &stack)| {
add(&sum, &compose(input, stack))
})
}
fn affine_compose_default<T>(
input: &T,
input_scale: f64,
input_shift: f64,
derivative_stack: [f64; 5],
scale: impl Fn(&T, f64) -> T,
add_constant: impl Fn(&T, f64) -> T,
compose: impl Fn(&T, [f64; 5]) -> T,
) -> T {
compose(
&add_constant(&scale(input, input_scale), input_shift),
derivative_stack,
)
}
fn affine_composed_sum_default<T>(
inputs: &[T],
input_scales: &[f64],
derivative_stacks: &[[f64; 5]],
constant: impl Fn(f64) -> T,
add: impl Fn(&T, &T) -> T,
scale: impl Fn(&T, f64) -> T,
add_constant: impl Fn(&T, f64) -> T,
compose: impl Fn(&T, [f64; 5]) -> T,
) -> T {
assert_eq!(inputs.len(), input_scales.len());
assert_eq!(inputs.len(), derivative_stacks.len());
inputs.iter().zip(input_scales).zip(derivative_stacks).fold(
constant(0.0),
|sum, ((input, &input_scale), &stack)| {
add(
&sum,
&affine_compose_default(
input,
input_scale,
0.0,
stack,
&scale,
&add_constant,
&compose,
),
)
},
)
}
fn shared_multiply_add_affine_composed_sum_default<T, const N: usize>(
lefts: &[&T; N],
right: &T,
addend: &T,
addend_scales: &[f64; N],
input_scales: &[f64; N],
derivative_stacks: &[[f64; 5]; N],
constant: impl Fn(f64) -> T,
add: impl Fn(&T, &T) -> T,
mul: impl Fn(&T, &T) -> T,
scale: impl Fn(&T, f64) -> T,
multiply_add: impl Fn(&T, &T, &T) -> T,
affine_compose: impl Fn(&T, f64, f64, [f64; 5]) -> T,
) -> T {
let (representatives, term_sources, source_count) =
canonical_shared_source_schedule(|term, representative| {
std::ptr::eq(lefts[term], lefts[representative])
&& addend_scales[term] == addend_scales[representative]
});
let (value, source_derivatives) =
aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
(0..source_count).fold(constant(value), |sum, source| {
let term = representatives[source];
let inner = if addend_scales[term] == 0.0 {
mul(lefts[term], right)
} else if addend_scales[term] == 1.0 {
multiply_add(lefts[term], right, addend)
} else {
multiply_add(lefts[term], right, &scale(addend, addend_scales[term]))
};
let composed = affine_compose(&inner, 1.0, 0.0, source_derivatives[source]);
add(&sum, &composed)
})
}
#[inline(always)]
pub(crate) fn canonical_shared_source_schedule<const N: usize>(
mut equivalent: impl FnMut(usize, usize) -> bool,
) -> ([usize; N], [usize; N], usize) {
let mut representatives = [0; N];
let mut term_sources = [0; N];
let mut source_count = 0;
for term in 0..N {
let mut source = 0;
while source < source_count && !equivalent(term, representatives[source]) {
source += 1;
}
if source == source_count {
representatives[source] = term;
source_count += 1;
}
term_sources[term] = source;
}
(representatives, term_sources, source_count)
}
#[inline(always)]
pub(crate) fn aggregate_shared_source_derivatives<const N: usize>(
term_sources: &[usize; N],
input_scales: &[f64; N],
derivative_stacks: &[[f64; 5]; N],
) -> (f64, [[f64; 5]; N]) {
let mut value = 0.0;
let mut source_derivatives = [[0.0; 5]; N];
for term in 0..N {
value += derivative_stacks[term][0];
let source = term_sources[term];
let mut scale_power = input_scales[term];
for order in 1..5 {
source_derivatives[source][order] += derivative_stacks[term][order] * scale_power;
scale_power *= input_scales[term];
}
}
(value, source_derivatives)
}
pub trait JetScalar<const K: usize>: crate::nested_dual::JetField + Copy {
fn constant(c: f64) -> Self;
fn variable(x: f64, axis: usize) -> Self;
fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
inputs: &[Self],
coefficients: &C,
) -> Self {
symmetric_quadratic_form_default(
inputs,
coefficients,
Self::constant,
crate::nested_dual::JetField::add,
crate::nested_dual::JetField::mul,
crate::nested_dual::JetField::scale,
)
}
fn linear_combination(inputs: &[Self], weights: &[f64]) -> Self {
linear_combination_default(
inputs,
weights,
Self::constant,
crate::nested_dual::JetField::add,
crate::nested_dual::JetField::scale,
)
}
fn add_constant(&self, constant: f64) -> Self {
self.add(&Self::constant(constant))
}
fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
multiply_add_default(
self,
right,
addend,
crate::nested_dual::JetField::mul,
crate::nested_dual::JetField::add,
)
}
fn composed_sum(inputs: &[Self], derivative_stacks: &[[f64; 5]]) -> Self {
composed_sum_default(
inputs,
derivative_stacks,
Self::constant,
crate::nested_dual::JetField::add,
crate::nested_dual::JetField::compose_unary,
)
}
fn product(&self, right: &Self) -> Self {
self.mul(right)
}
fn affine_compose(
&self,
input_scale: f64,
input_shift: f64,
derivative_stack: [f64; 5],
) -> Self {
affine_compose_default(
self,
input_scale,
input_shift,
derivative_stack,
crate::nested_dual::JetField::scale,
Self::add_constant,
crate::nested_dual::JetField::compose_unary,
)
}
fn affine_composed_sum(
inputs: &[Self],
input_scales: &[f64],
derivative_stacks: &[[f64; 5]],
) -> Self {
affine_composed_sum_default(
inputs,
input_scales,
derivative_stacks,
Self::constant,
crate::nested_dual::JetField::add,
crate::nested_dual::JetField::scale,
Self::add_constant,
crate::nested_dual::JetField::compose_unary,
)
}
fn shared_multiply_add_affine_composed_sum<const N: usize>(
lefts: &[&Self; N],
right: &Self,
addend: &Self,
addend_scales: &[f64; N],
input_scales: &[f64; N],
derivative_stacks: &[[f64; 5]; N],
) -> Self {
shared_multiply_add_affine_composed_sum_default(
lefts,
right,
addend,
addend_scales,
input_scales,
derivative_stacks,
Self::constant,
crate::nested_dual::JetField::add,
crate::nested_dual::JetField::mul,
crate::nested_dual::JetField::scale,
Self::multiply_add,
Self::affine_compose,
)
}
fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 5]) -> Self {
self.compose_unary(stack_fn(self.value()))
}
fn exp(&self) -> Self {
let e = self.value().exp();
self.compose_unary([e, e, e, e, e])
}
fn sqrt(&self) -> Self {
let u = self.value();
let s = u.sqrt();
self.compose_unary([
s,
0.5 / s,
-0.25 / (u * s),
0.375 / (u * u * s),
-0.9375 / (u * u * u * s),
])
}
fn ln(&self) -> Self {
let u = self.value();
let r = 1.0 / u;
self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
}
fn recip(&self) -> Self {
let r = 1.0 / self.value();
let r2 = r * r;
self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
}
fn powf(&self, a: f64) -> Self {
let u = self.value();
self.compose_unary([
u.powf(a),
a * u.powf(a - 1.0),
a * (a - 1.0) * u.powf(a - 2.0),
a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
])
}
fn ln_gamma(&self) -> Self {
self.compose_unary(crate::jet_tower::ln_gamma_derivative_stack(self.value()))
}
fn digamma(&self) -> Self {
self.compose_unary(crate::jet_tower::digamma_derivative_stack(self.value()))
}
}
impl<S, const K: usize> JetScalar<K> for crate::nested_dual::Dual2<S>
where
S: JetScalar<K>,
{
#[inline]
fn constant(c: f64) -> Self {
Self {
v: S::constant(c),
g: S::constant(0.0),
h: S::constant(0.0),
}
}
#[inline]
fn variable(x: f64, axis: usize) -> Self {
Self {
v: S::variable(x, axis),
g: S::constant(0.0),
h: S::constant(0.0),
}
}
}
pub trait RuntimeJetScalar<'arena>: Clone {
type Workspace: ?Sized;
fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
fn constant_like(&self, c: f64) -> Self;
fn with_value(&self, value: f64) -> Self;
fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
inputs: &[Self],
coefficients: &C,
dimension: usize,
workspace: &'arena Self::Workspace,
) -> Self {
symmetric_quadratic_form_default(
inputs,
coefficients,
|value| Self::constant(value, dimension, workspace),
Self::add,
Self::mul,
Self::scale,
)
}
fn linear_combination(
inputs: &[Self],
weights: &[f64],
dimension: usize,
workspace: &'arena Self::Workspace,
) -> Self {
linear_combination_default(
inputs,
weights,
|value| Self::constant(value, dimension, workspace),
Self::add,
Self::scale,
)
}
fn add_constant(&self, constant: f64) -> Self {
self.with_value(self.value() + constant)
}
fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
multiply_add_default(self, right, addend, Self::mul, Self::add)
}
fn composed_sum(
inputs: &[Self],
derivative_stacks: &[[f64; 5]],
dimension: usize,
workspace: &'arena Self::Workspace,
) -> Self {
composed_sum_default(
inputs,
derivative_stacks,
|value| Self::constant(value, dimension, workspace),
Self::add,
Self::compose_unary,
)
}
fn product(&self, right: &Self) -> Self {
self.mul(right)
}
fn affine_compose(
&self,
input_scale: f64,
input_shift: f64,
derivative_stack: [f64; 5],
) -> Self {
affine_compose_default(
self,
input_scale,
input_shift,
derivative_stack,
Self::scale,
|value, constant| value.add_constant(constant),
Self::compose_unary,
)
}
fn affine_composed_sum(
inputs: &[Self],
input_scales: &[f64],
derivative_stacks: &[[f64; 5]],
dimension: usize,
workspace: &'arena Self::Workspace,
) -> Self {
affine_composed_sum_default(
inputs,
input_scales,
derivative_stacks,
|value| Self::constant(value, dimension, workspace),
Self::add,
Self::scale,
|value, constant| value.add_constant(constant),
Self::compose_unary,
)
}
fn shared_multiply_add_affine_composed_sum<const N: usize>(
lefts: &[&Self; N],
right: &Self,
addend: &Self,
addend_scales: &[f64; N],
input_scales: &[f64; N],
derivative_stacks: &[[f64; 5]; N],
dimension: usize,
workspace: &'arena Self::Workspace,
) -> Self {
shared_multiply_add_affine_composed_sum_default(
lefts,
right,
addend,
addend_scales,
input_scales,
derivative_stacks,
|value| Self::constant(value, dimension, workspace),
Self::add,
Self::mul,
Self::scale,
Self::multiply_add,
|input, scale, shift, stack| input.affine_compose(scale, shift, stack),
)
}
fn dimension(&self) -> usize;
fn value(&self) -> f64;
fn add(&self, o: &Self) -> Self;
fn sub(&self, o: &Self) -> Self;
fn mul(&self, o: &Self) -> Self;
fn neg(&self) -> Self;
fn scale(&self, s: f64) -> Self;
fn compose_unary(&self, d: [f64; 5]) -> Self;
fn exp(&self) -> Self {
let e = self.value().exp();
self.compose_unary([e, e, e, e, e])
}
fn ln(&self) -> Self {
let u = self.value();
let r = 1.0 / u;
self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
}
fn recip(&self) -> Self {
let r = 1.0 / self.value();
let r2 = r * r;
self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct RuntimeValue {
value: f64,
dimension: usize,
}
impl<'arena> RuntimeJetScalar<'arena> for RuntimeValue {
type Workspace = ();
#[inline(always)]
fn constant(c: f64, dimension: usize, &(): &'arena Self::Workspace) -> Self {
Self {
value: c,
dimension,
}
}
#[inline(always)]
fn variable(x: f64, axis: usize, dimension: usize, &(): &'arena Self::Workspace) -> Self {
assert!(
axis < dimension,
"runtime value variable axis out of bounds"
);
Self {
value: x,
dimension,
}
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
Self {
value: c,
dimension: self.dimension,
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
value,
dimension: self.dimension,
}
}
#[inline(always)]
fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
inputs: &[Self],
coefficients: &C,
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(inputs.len(), coefficients.dimension());
assert!(inputs.iter().all(|input| input.dimension == dimension));
Self {
value: coefficients.quadratic_value(inputs, |input| input.value),
dimension,
}
}
#[inline(always)]
fn linear_combination(
inputs: &[Self],
weights: &[f64],
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(inputs.len(), weights.len());
assert!(inputs.iter().all(|input| input.dimension == dimension));
let value = inputs
.iter()
.zip(weights)
.map(|(input, &weight)| input.value * weight)
.sum();
Self { value, dimension }
}
#[inline(always)]
fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
self.assert_same_dimension(right);
self.assert_same_dimension(addend);
Self {
value: self.value * right.value + addend.value,
dimension: self.dimension,
}
}
#[inline(always)]
fn composed_sum(
inputs: &[Self],
derivative_stacks: &[[f64; 5]],
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(inputs.len(), derivative_stacks.len());
assert!(inputs.iter().all(|input| input.dimension == dimension));
Self {
value: derivative_stacks.iter().map(|stack| stack[0]).sum(),
dimension,
}
}
#[inline(always)]
fn product(&self, right: &Self) -> Self {
self.mul(right)
}
#[inline(always)]
fn affine_compose(
&self,
input_scale: f64,
input_shift: f64,
derivative_stack: [f64; 5],
) -> Self {
affine_compose_default(
self,
input_scale,
input_shift,
derivative_stack,
Self::scale,
|value, constant| value.add_constant(constant),
Self::compose_unary,
)
}
#[inline(always)]
fn affine_composed_sum(
inputs: &[Self],
input_scales: &[f64],
derivative_stacks: &[[f64; 5]],
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(inputs.len(), input_scales.len());
assert_eq!(inputs.len(), derivative_stacks.len());
assert!(inputs.iter().all(|input| input.dimension == dimension));
Self {
value: derivative_stacks.iter().map(|stack| stack[0]).sum(),
dimension,
}
}
#[inline(always)]
fn dimension(&self) -> usize {
self.dimension
}
#[inline(always)]
fn value(&self) -> f64 {
self.value
}
#[inline(always)]
fn add(&self, other: &Self) -> Self {
self.assert_same_dimension(other);
Self {
value: self.value + other.value,
dimension: self.dimension,
}
}
#[inline(always)]
fn sub(&self, other: &Self) -> Self {
self.assert_same_dimension(other);
Self {
value: self.value - other.value,
dimension: self.dimension,
}
}
#[inline(always)]
fn mul(&self, other: &Self) -> Self {
self.assert_same_dimension(other);
Self {
value: self.value * other.value,
dimension: self.dimension,
}
}
#[inline(always)]
fn neg(&self) -> Self {
Self {
value: -self.value,
dimension: self.dimension,
}
}
#[inline(always)]
fn scale(&self, scale: f64) -> Self {
Self {
value: self.value * scale,
dimension: self.dimension,
}
}
#[inline(always)]
fn compose_unary(&self, derivative_stack: [f64; 5]) -> Self {
Self {
value: derivative_stack[0],
dimension: self.dimension,
}
}
}
impl RuntimeValue {
#[inline(always)]
fn assert_same_dimension(&self, other: &Self) {
assert_eq!(self.dimension, other.dimension);
}
}
#[derive(Clone, Copy, Debug)]
#[repr(transparent)]
pub struct FixedRuntimeJet<S, const K: usize> {
inner: S,
}
impl<S, const K: usize> FixedRuntimeJet<S, K> {
#[inline(always)]
#[must_use]
pub fn from_inner(inner: S) -> Self {
Self { inner }
}
#[inline(always)]
#[must_use]
pub fn into_inner(self) -> S {
self.inner
}
}
impl<'arena, S: JetScalar<K>, const K: usize> RuntimeJetScalar<'arena> for FixedRuntimeJet<S, K> {
type Workspace = ();
#[inline(always)]
fn constant(c: f64, dimension: usize, &(): &'arena Self::Workspace) -> Self {
assert_eq!(dimension, K, "fixed jet dimension mismatch");
Self {
inner: S::constant(c),
}
}
#[inline(always)]
fn variable(x: f64, axis: usize, dimension: usize, &(): &'arena Self::Workspace) -> Self {
assert_eq!(dimension, K, "fixed jet dimension mismatch");
Self {
inner: S::variable(x, axis),
}
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
Self {
inner: S::constant(c),
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
inner: self.inner.compose_unary([value, 1.0, 0.0, 0.0, 0.0]),
}
}
#[inline(always)]
fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
inputs: &[Self],
coefficients: &C,
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(dimension, K, "fixed jet dimension mismatch");
assert_eq!(inputs.len(), coefficients.dimension());
let inner =
unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
Self {
inner: S::symmetric_quadratic_form(inner, coefficients),
}
}
#[inline(always)]
fn linear_combination(
inputs: &[Self],
weights: &[f64],
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(dimension, K, "fixed jet dimension mismatch");
assert_eq!(inputs.len(), weights.len());
let inner =
unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
Self {
inner: S::linear_combination(inner, weights),
}
}
#[inline(always)]
fn add_constant(&self, constant: f64) -> Self {
Self {
inner: self.inner.add_constant(constant),
}
}
#[inline(always)]
fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
Self {
inner: self.inner.multiply_add(&right.inner, &addend.inner),
}
}
#[inline(always)]
fn composed_sum(
inputs: &[Self],
derivative_stacks: &[[f64; 5]],
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(dimension, K, "fixed jet dimension mismatch");
let inner =
unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
Self {
inner: S::composed_sum(inner, derivative_stacks),
}
}
#[inline(always)]
fn product(&self, right: &Self) -> Self {
Self {
inner: self.inner.product(&right.inner),
}
}
#[inline(always)]
fn affine_compose(
&self,
input_scale: f64,
input_shift: f64,
derivative_stack: [f64; 5],
) -> Self {
Self {
inner: self
.inner
.affine_compose(input_scale, input_shift, derivative_stack),
}
}
#[inline(always)]
fn affine_composed_sum(
inputs: &[Self],
input_scales: &[f64],
derivative_stacks: &[[f64; 5]],
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(dimension, K, "fixed jet dimension mismatch");
let inner =
unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
Self {
inner: S::affine_composed_sum(inner, input_scales, derivative_stacks),
}
}
#[inline(always)]
fn shared_multiply_add_affine_composed_sum<const N: usize>(
lefts: &[&Self; N],
right: &Self,
addend: &Self,
addend_scales: &[f64; N],
input_scales: &[f64; N],
derivative_stacks: &[[f64; 5]; N],
dimension: usize,
&(): &'arena Self::Workspace,
) -> Self {
assert_eq!(dimension, K, "fixed jet dimension mismatch");
let left_inner: [&S; N] = std::array::from_fn(|term| &lefts[term].inner);
Self {
inner: S::shared_multiply_add_affine_composed_sum(
&left_inner,
&right.inner,
&addend.inner,
addend_scales,
input_scales,
derivative_stacks,
),
}
}
#[inline(always)]
fn dimension(&self) -> usize {
K
}
#[inline(always)]
fn value(&self) -> f64 {
self.inner.value()
}
#[inline(always)]
fn add(&self, o: &Self) -> Self {
Self {
inner: self.inner.add(&o.inner),
}
}
#[inline(always)]
fn sub(&self, o: &Self) -> Self {
Self {
inner: self.inner.sub(&o.inner),
}
}
#[inline(always)]
fn mul(&self, o: &Self) -> Self {
Self {
inner: self.inner.mul(&o.inner),
}
}
#[inline(always)]
fn neg(&self) -> Self {
Self {
inner: self.inner.neg(),
}
}
#[inline(always)]
fn scale(&self, s: f64) -> Self {
Self {
inner: self.inner.scale(s),
}
}
#[inline(always)]
fn compose_unary(&self, d: [f64; 5]) -> Self {
Self {
inner: self.inner.compose_unary(d),
}
}
}
#[derive(Debug)]
pub struct DynamicJetArena {
bump: bumpalo::Bump,
}
impl DynamicJetArena {
#[must_use]
pub fn new() -> Self {
Self {
bump: bumpalo::Bump::new(),
}
}
#[must_use]
pub fn with_capacity(bytes: usize) -> Self {
Self {
bump: bumpalo::Bump::with_capacity(bytes),
}
}
pub fn reset(&mut self) {
let high_water = self.bump.allocated_bytes();
self.bump.reset();
if self.bump.allocated_bytes() < high_water {
self.bump = bumpalo::Bump::with_capacity(high_water);
}
}
#[must_use]
pub fn allocated_bytes(&self) -> usize {
self.bump.allocated_bytes()
}
#[inline(always)]
fn zeros(&self, len: usize) -> &mut [f64] {
self.bump.alloc_slice_fill_copy(len, 0.0)
}
#[inline(always)]
pub fn alloc_slice_fill_with<T>(&self, len: usize, fill: impl FnMut(usize) -> T) -> &mut [T] {
self.bump.alloc_slice_fill_with(len, fill)
}
}
impl Default for DynamicJetArena {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Copy, Debug)]
pub struct DynamicOrder1<'arena> {
arena: &'arena DynamicJetArena,
pub v: f64,
pub g: &'arena [f64],
}
impl DynamicOrder1<'_> {
#[inline]
#[must_use]
pub fn g(&self) -> &[f64] {
self.g
}
#[inline]
fn assert_compatible(&self, o: &Self) {
assert_eq!(
self.g.len(),
o.g.len(),
"dynamic first-order jet dimension mismatch"
);
assert!(
std::ptr::eq(self.arena, o.arena),
"dynamic jets belong to different arenas"
);
}
}
impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder1<'arena> {
type Workspace = DynamicJetArena;
fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
Self {
arena,
v: c,
g: arena.zeros(dimension),
}
}
fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
assert!(
axis < dimension,
"dynamic first-order jet axis out of bounds"
);
let g = arena.zeros(dimension);
g[axis] = 1.0;
Self { arena, v: x, g }
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
Self {
arena: self.arena,
v: c,
g: self.arena.zeros(self.dimension()),
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
arena: self.arena,
v: value,
g: self.g,
}
}
fn dimension(&self) -> usize {
self.g.len()
}
fn value(&self) -> f64 {
self.v
}
fn add(&self, o: &Self) -> Self {
self.assert_compatible(o);
let g = self.arena.zeros(self.dimension());
for i in 0..g.len() {
g[i] = self.g[i] + o.g[i];
}
Self {
arena: self.arena,
v: self.v + o.v,
g,
}
}
fn sub(&self, o: &Self) -> Self {
self.assert_compatible(o);
let g = self.arena.zeros(self.dimension());
for i in 0..g.len() {
g[i] = self.g[i] - o.g[i];
}
Self {
arena: self.arena,
v: self.v - o.v,
g,
}
}
fn mul(&self, o: &Self) -> Self {
self.assert_compatible(o);
let g = self.arena.zeros(self.dimension());
for i in 0..g.len() {
g[i] = self.v * o.g[i] + self.g[i] * o.v;
}
Self {
arena: self.arena,
v: self.v * o.v,
g,
}
}
fn neg(&self) -> Self {
self.scale(-1.0)
}
fn scale(&self, s: f64) -> Self {
let g = self.arena.zeros(self.dimension());
for i in 0..g.len() {
g[i] = self.g[i] * s;
}
Self {
arena: self.arena,
v: self.v * s,
g,
}
}
fn compose_unary(&self, d: [f64; 5]) -> Self {
let g = self.arena.zeros(self.dimension());
for i in 0..g.len() {
g[i] = d[1] * self.g[i];
}
Self {
arena: self.arena,
v: d[0],
g,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct DynamicOrder2<'arena> {
arena: &'arena DynamicJetArena,
pub v: f64,
pub g: &'arena [f64],
pub h: &'arena [f64],
}
impl DynamicOrder2<'_> {
#[inline]
#[must_use]
pub fn from_channel_functions<'arena>(
value: f64,
dimension: usize,
arena: &'arena DynamicJetArena,
mut gradient: impl FnMut(usize) -> f64,
mut hessian: impl FnMut(usize, usize) -> f64,
) -> DynamicOrder2<'arena> {
let g = arena.alloc_slice_fill_with(dimension, |axis| gradient(axis));
let h = arena.zeros(dimension * dimension);
for row in 0..dimension {
for column in row..dimension {
let channel = hessian(row, column);
h[row * dimension + column] = channel;
h[column * dimension + row] = channel;
}
}
DynamicOrder2 {
arena,
v: value,
g,
h,
}
}
#[inline]
#[must_use]
pub fn g(&self) -> &[f64] {
self.g
}
#[inline]
#[must_use]
pub fn h(&self) -> &[f64] {
self.h
}
#[inline]
#[must_use]
pub fn h_at(&self, row: usize, col: usize) -> f64 {
self.h[row * self.dimension() + col]
}
#[inline(always)]
fn assert_compatible(&self, o: &Self) {
assert_eq!(
self.g.len(),
o.g.len(),
"dynamic second-order jet dimension mismatch"
);
assert_eq!(
self.h.len(),
o.h.len(),
"dynamic second-order jet Hessian mismatch"
);
assert!(
std::ptr::eq(self.arena, o.arena),
"dynamic jets belong to different arenas"
);
}
}
impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder2<'arena> {
type Workspace = DynamicJetArena;
#[inline(always)]
fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
Self {
arena,
v: c,
g: arena.zeros(dimension),
h: arena.zeros(dimension * dimension),
}
}
#[inline(always)]
fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
assert!(
axis < dimension,
"dynamic second-order jet axis out of bounds"
);
let g = arena.zeros(dimension);
g[axis] = 1.0;
Self {
arena,
v: x,
g,
h: arena.zeros(dimension * dimension),
}
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
let dimension = self.dimension();
Self {
arena: self.arena,
v: c,
g: self.arena.zeros(dimension),
h: self.arena.zeros(dimension * dimension),
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
arena: self.arena,
v: value,
g: self.g,
h: self.h,
}
}
#[inline(always)]
fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
inputs: &[Self],
coefficients: &C,
dimension: usize,
arena: &'arena DynamicJetArena,
) -> Self {
assert_eq!(inputs.len(), coefficients.dimension());
assert!(
inputs.iter().all(|input| {
input.dimension() == dimension && std::ptr::eq(input.arena, arena)
}),
"dynamic quadratic-form jets must share dimension and arena"
);
let input_dimension = inputs.len();
let values = arena.zeros(input_dimension);
for (value, input) in values.iter_mut().zip(inputs) {
*value = input.v;
}
let projected = arena.zeros(input_dimension);
coefficients.multiply(values, projected);
let mut value = 0.0;
for axis in 0..input_dimension {
value += values[axis] * projected[axis];
}
let gradient = arena.zeros(dimension);
for primary in 0..dimension {
let mut channel = 0.0;
for axis in 0..input_dimension {
channel += projected[axis] * inputs[axis].g[primary];
}
gradient[primary] = 2.0 * channel;
}
let hessian = arena.zeros(dimension * dimension);
let input_gradient = arena.zeros(input_dimension);
let projected_gradient = arena.zeros(input_dimension);
for primary_b in 0..dimension {
for row in 0..input_dimension {
input_gradient[row] = inputs[row].g[primary_b];
}
coefficients.multiply(input_gradient, projected_gradient);
for primary_a in 0..=primary_b {
let mut inherited = 0.0;
let mut curvature = 0.0;
for row in 0..input_dimension {
inherited += projected[row] * inputs[row].h[primary_a * dimension + primary_b];
curvature += inputs[row].g[primary_a] * projected_gradient[row];
}
let channel = 2.0 * (inherited + curvature);
hessian[primary_a * dimension + primary_b] = channel;
hessian[primary_b * dimension + primary_a] = channel;
}
}
Self {
arena,
v: value,
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn product(&self, right: &Self) -> Self {
self.assert_compatible(right);
let dimension = self.dimension();
let gradient = self.arena.zeros(dimension);
let hessian = self.arena.zeros(dimension * dimension);
for primary in 0..dimension {
gradient[primary] = self.v * right.g[primary] + self.g[primary] * right.v;
for other in primary..dimension {
let index = primary * dimension + other;
let channel = self.v * right.h[index]
+ self.g[primary] * right.g[other]
+ self.g[other] * right.g[primary]
+ self.h[index] * right.v;
hessian[index] = channel;
hessian[other * dimension + primary] = channel;
}
}
Self {
arena: self.arena,
v: self.v * right.v,
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn affine_compose(
&self,
input_scale: f64,
input_shift: f64,
derivative_stack: [f64; 5],
) -> Self {
assert!(input_shift.is_finite(), "affine input shift must be finite");
let arena = self.arena;
let dimension = self.dimension();
let first = derivative_stack[1] * input_scale;
let second = derivative_stack[2] * input_scale * input_scale;
let gradient = arena.zeros(dimension);
let hessian = arena.zeros(dimension * dimension);
for primary in 0..dimension {
gradient[primary] = first * self.g[primary];
for other in primary..dimension {
let index = primary * dimension + other;
let channel = first * self.h[index] + second * self.g[primary] * self.g[other];
hessian[index] = channel;
hessian[other * dimension + primary] = channel;
}
}
Self {
arena,
v: derivative_stack[0],
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn affine_composed_sum(
inputs: &[Self],
input_scales: &[f64],
derivative_stacks: &[[f64; 5]],
dimension: usize,
arena: &'arena DynamicJetArena,
) -> Self {
assert_eq!(inputs.len(), input_scales.len());
assert_eq!(inputs.len(), derivative_stacks.len());
assert!(
inputs.iter().all(|input| {
input.dimension() == dimension && std::ptr::eq(input.arena, arena)
}),
"dynamic affine-composed-sum jets must share dimension and arena"
);
let gradient = arena.zeros(dimension);
let hessian = arena.zeros(dimension * dimension);
let mut value = 0.0;
for ((input, &input_scale), stack) in inputs.iter().zip(input_scales).zip(derivative_stacks)
{
let first = stack[1] * input_scale;
let second = stack[2] * input_scale * input_scale;
value += stack[0];
for primary in 0..dimension {
gradient[primary] += first * input.g[primary];
for other in primary..dimension {
let index = primary * dimension + other;
hessian[index] +=
first * input.h[index] + second * input.g[primary] * input.g[other];
}
}
}
for primary in 0..dimension {
for other in primary + 1..dimension {
hessian[other * dimension + primary] = hessian[primary * dimension + other];
}
}
Self {
arena,
v: value,
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn shared_multiply_add_affine_composed_sum<const N: usize>(
lefts: &[&Self; N],
right: &Self,
addend: &Self,
addend_scales: &[f64; N],
input_scales: &[f64; N],
derivative_stacks: &[[f64; 5]; N],
dimension: usize,
arena: &'arena DynamicJetArena,
) -> Self {
assert!(
lefts.iter().all(|input| {
input.dimension() == dimension && std::ptr::eq(input.arena, arena)
}) && (N == 0 || (right.dimension() == dimension && std::ptr::eq(right.arena, arena))),
"dynamic fused product-composition jets must share dimension and arena"
);
let addend_live = addend_scales.iter().any(|&scale| scale != 0.0);
assert!(
!addend_live || (addend.dimension() == dimension && std::ptr::eq(addend.arena, arena)),
"live dynamic fused addends must share dimension and arena"
);
let (representatives, term_sources, source_count) =
canonical_shared_source_schedule::<N>(|term, representative| {
std::ptr::eq(lefts[term], lefts[representative])
&& addend_scales[term] == addend_scales[representative]
});
let (value, source_derivatives) =
aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
let source_gradients = arena.zeros(source_count * dimension);
let gradient = arena.zeros(dimension);
let hessian = arena.zeros(dimension * dimension);
let mut right_first = 0.0;
let mut addend_first = 0.0;
for source in 0..source_count {
let term = representatives[source];
let first = source_derivatives[source][1];
right_first += first * lefts[term].v;
addend_first += first * addend_scales[term];
for primary in 0..dimension {
let product_gradient =
lefts[term].v * right.g[primary] + lefts[term].g[primary] * right.v;
let inner_gradient = if addend_scales[term] == 0.0 {
product_gradient
} else if addend_scales[term] == 1.0 {
product_gradient + addend.g[primary]
} else {
product_gradient + addend_scales[term] * addend.g[primary]
};
source_gradients[source * dimension + primary] = inner_gradient;
gradient[primary] += first * lefts[term].g[primary] * right.v;
}
}
if N != 0 {
for primary in 0..dimension {
gradient[primary] += right_first * right.g[primary];
}
}
if addend_live {
for primary in 0..dimension {
gradient[primary] += addend_first * addend.g[primary];
}
}
for primary in 0..dimension {
for other in primary..dimension {
let index = primary * dimension + other;
let mut channel = if N == 0 {
0.0
} else {
right_first * right.h[index]
};
if addend_live {
channel += addend_first * addend.h[index];
}
for source in 0..source_count {
let term = representatives[source];
let local_product_hessian = lefts[term].g[primary] * right.g[other]
+ lefts[term].g[other] * right.g[primary]
+ lefts[term].h[index] * right.v;
channel += source_derivatives[source][1] * local_product_hessian
+ source_derivatives[source][2]
* source_gradients[source * dimension + primary]
* source_gradients[source * dimension + other];
}
hessian[index] = channel;
hessian[other * dimension + primary] = channel;
}
}
Self {
arena,
v: value,
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
self.assert_compatible(right);
self.assert_compatible(addend);
let dimension = self.dimension();
let gradient = self.arena.zeros(dimension);
let hessian = self.arena.zeros(dimension * dimension);
for primary in 0..dimension {
gradient[primary] =
self.v * right.g[primary] + self.g[primary] * right.v + addend.g[primary];
for other in primary..dimension {
let index = primary * dimension + other;
let channel = self.v * right.h[index]
+ self.g[primary] * right.g[other]
+ self.g[other] * right.g[primary]
+ self.h[index] * right.v
+ addend.h[index];
hessian[index] = channel;
hessian[other * dimension + primary] = channel;
}
}
Self {
arena: self.arena,
v: self.v * right.v + addend.v,
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn composed_sum(
inputs: &[Self],
derivative_stacks: &[[f64; 5]],
dimension: usize,
arena: &'arena DynamicJetArena,
) -> Self {
assert_eq!(inputs.len(), derivative_stacks.len());
assert!(
inputs.iter().all(|input| {
input.dimension() == dimension && std::ptr::eq(input.arena, arena)
}),
"dynamic composed-sum jets must share dimension and arena"
);
let gradient = arena.zeros(dimension);
let hessian = arena.zeros(dimension * dimension);
let mut value = 0.0;
for (input, stack) in inputs.iter().zip(derivative_stacks) {
value += stack[0];
for primary in 0..dimension {
gradient[primary] += stack[1] * input.g[primary];
for other in primary..dimension {
let index = primary * dimension + other;
hessian[index] +=
stack[1] * input.h[index] + stack[2] * input.g[primary] * input.g[other];
}
}
}
for primary in 0..dimension {
for other in primary + 1..dimension {
hessian[other * dimension + primary] = hessian[primary * dimension + other];
}
}
Self {
arena,
v: value,
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn linear_combination(
inputs: &[Self],
weights: &[f64],
dimension: usize,
arena: &'arena DynamicJetArena,
) -> Self {
assert_eq!(inputs.len(), weights.len());
assert!(
inputs.iter().all(|input| {
input.dimension() == dimension && std::ptr::eq(input.arena, arena)
}),
"dynamic linear-combination jets must share dimension and arena"
);
let mut value = 0.0;
for (input, &weight) in inputs.iter().zip(weights) {
value += input.v * weight;
}
let gradient = arena.zeros(dimension);
let hessian = arena.zeros(dimension * dimension);
for primary in 0..dimension {
for (input, &weight) in inputs.iter().zip(weights) {
gradient[primary] += input.g[primary] * weight;
}
for other in primary..dimension {
let index = primary * dimension + other;
for (input, &weight) in inputs.iter().zip(weights) {
hessian[index] += input.h[index] * weight;
}
hessian[other * dimension + primary] = hessian[index];
}
}
Self {
arena,
v: value,
g: gradient,
h: hessian,
}
}
#[inline(always)]
fn dimension(&self) -> usize {
self.g.len()
}
#[inline(always)]
fn value(&self) -> f64 {
self.v
}
#[inline(always)]
fn add(&self, o: &Self) -> Self {
self.assert_compatible(o);
let dimension = self.dimension();
let g = self.arena.zeros(dimension);
let h = self.arena.zeros(self.h.len());
for i in 0..g.len() {
g[i] = self.g[i] + o.g[i];
}
for row in 0..dimension {
for column in row..dimension {
let index = row * dimension + column;
let channel = self.h[index] + o.h[index];
h[index] = channel;
h[column * dimension + row] = channel;
}
}
Self {
arena: self.arena,
v: self.v + o.v,
g,
h,
}
}
#[inline(always)]
fn sub(&self, o: &Self) -> Self {
self.assert_compatible(o);
let dimension = self.dimension();
let g = self.arena.zeros(dimension);
let h = self.arena.zeros(self.h.len());
for i in 0..g.len() {
g[i] = self.g[i] - o.g[i];
}
for row in 0..dimension {
for column in row..dimension {
let index = row * dimension + column;
let channel = self.h[index] - o.h[index];
h[index] = channel;
h[column * dimension + row] = channel;
}
}
Self {
arena: self.arena,
v: self.v - o.v,
g,
h,
}
}
#[inline(always)]
fn mul(&self, o: &Self) -> Self {
self.assert_compatible(o);
let n = self.dimension();
let g = self.arena.zeros(n);
let h = self.arena.zeros(n * n);
for i in 0..n {
g[i] = self.v * o.g[i] + self.g[i] * o.v;
}
for i in 0..n {
for j in i..n {
let ij = i * n + j;
let hij =
self.v * o.h[ij] + self.g[i] * o.g[j] + self.g[j] * o.g[i] + self.h[ij] * o.v;
h[ij] = hij;
h[j * n + i] = hij;
}
}
Self {
arena: self.arena,
v: self.v * o.v,
g,
h,
}
}
#[inline(always)]
fn neg(&self) -> Self {
self.scale(-1.0)
}
#[inline(always)]
fn scale(&self, s: f64) -> Self {
let dimension = self.dimension();
let g = self.arena.zeros(dimension);
let h = self.arena.zeros(self.h.len());
for i in 0..g.len() {
g[i] = self.g[i] * s;
}
for row in 0..dimension {
for column in row..dimension {
let index = row * dimension + column;
let channel = self.h[index] * s;
h[index] = channel;
h[column * dimension + row] = channel;
}
}
Self {
arena: self.arena,
v: self.v * s,
g,
h,
}
}
#[inline(always)]
fn compose_unary(&self, d: [f64; 5]) -> Self {
let n = self.dimension();
let g = self.arena.zeros(n);
let h = self.arena.zeros(n * n);
for i in 0..n {
g[i] = d[1] * self.g[i];
}
for i in 0..n {
for j in i..n {
let ij = i * n + j;
let channel = d[1] * self.h[ij] + d[2] * self.g[i] * self.g[j];
h[ij] = channel;
h[j * n + i] = channel;
}
}
Self {
arena: self.arena,
v: d[0],
g,
h,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct DynamicOneSeed<'arena> {
pub base: DynamicOrder2<'arena>,
pub eps: DynamicOrder2<'arena>,
}
impl<'arena> DynamicOneSeed<'arena> {
#[inline(always)]
#[must_use]
pub fn seed_direction(
x: f64,
axis: usize,
u_axis: f64,
dimension: usize,
arena: &'arena DynamicJetArena,
) -> Self {
Self {
base: DynamicOrder2::variable(x, axis, dimension, arena),
eps: DynamicOrder2::constant(u_axis, dimension, arena),
}
}
#[inline(always)]
#[must_use]
pub fn contracted_third(&self) -> &[f64] {
self.eps.h()
}
}
impl<'arena> RuntimeJetScalar<'arena> for DynamicOneSeed<'arena> {
type Workspace = DynamicJetArena;
#[inline(always)]
fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
Self {
base: DynamicOrder2::constant(c, dimension, arena),
eps: DynamicOrder2::constant(0.0, dimension, arena),
}
}
#[inline(always)]
fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
Self {
base: DynamicOrder2::variable(x, axis, dimension, arena),
eps: DynamicOrder2::constant(0.0, dimension, arena),
}
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
Self {
base: self.base.constant_like(c),
eps: self.eps.constant_like(0.0),
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
base: self.base.with_value(value),
eps: self.eps,
}
}
#[inline(always)]
fn dimension(&self) -> usize {
self.base.dimension()
}
#[inline(always)]
fn value(&self) -> f64 {
self.base.value()
}
#[inline(always)]
fn add(&self, o: &Self) -> Self {
Self {
base: self.base.add(&o.base),
eps: self.eps.add(&o.eps),
}
}
#[inline(always)]
fn sub(&self, o: &Self) -> Self {
Self {
base: self.base.sub(&o.base),
eps: self.eps.sub(&o.eps),
}
}
#[inline(always)]
fn mul(&self, o: &Self) -> Self {
self.base.assert_compatible(&o.base);
self.eps.assert_compatible(&o.eps);
Self {
base: self.base.mul(&o.base),
eps: DynamicOrder2::from_channel_functions(
self.base.v * o.eps.v + self.eps.v * o.base.v,
self.dimension(),
self.base.arena,
|i| {
self.base.v * o.eps.g[i]
+ self.base.g[i] * o.eps.v
+ self.eps.v * o.base.g[i]
+ self.eps.g[i] * o.base.v
},
|i, j| {
let ij = i * self.dimension() + j;
self.base.v * o.eps.h[ij]
+ self.base.g[i] * o.eps.g[j]
+ self.base.g[j] * o.eps.g[i]
+ self.base.h[ij] * o.eps.v
+ self.eps.v * o.base.h[ij]
+ self.eps.g[i] * o.base.g[j]
+ self.eps.g[j] * o.base.g[i]
+ self.eps.h[ij] * o.base.v
},
),
}
}
#[inline(always)]
fn neg(&self) -> Self {
Self {
base: self.base.neg(),
eps: self.eps.neg(),
}
}
#[inline(always)]
fn scale(&self, s: f64) -> Self {
Self {
base: self.base.scale(s),
eps: self.eps.scale(s),
}
}
#[inline(always)]
fn compose_unary(&self, d: [f64; 5]) -> Self {
let base = self.base.compose_unary(d);
let dimension = self.dimension();
let eps = DynamicOrder2::from_channel_functions(
d[1] * self.eps.v,
dimension,
self.base.arena,
|i| d[2] * self.base.g[i] * self.eps.v + d[1] * self.eps.g[i],
|i, j| {
let ij = i * dimension + j;
d[1] * self.eps.h[ij]
+ d[2]
* (self.base.g[i] * self.eps.g[j]
+ self.base.g[j] * self.eps.g[i]
+ self.base.h[ij] * self.eps.v)
+ d[3] * self.base.g[i] * self.base.g[j] * self.eps.v
},
);
Self { base, eps }
}
}
#[derive(Debug)]
pub struct DynamicJetBatchWorkspace {
arena: DynamicJetArena,
lanes: usize,
}
impl DynamicJetBatchWorkspace {
#[must_use]
pub fn new(lanes: usize) -> Self {
Self {
arena: DynamicJetArena::new(),
lanes,
}
}
pub fn reset(&mut self, lanes: usize) {
self.arena.reset();
self.lanes = lanes;
}
#[must_use]
pub fn allocated_bytes(&self) -> usize {
self.arena.allocated_bytes()
}
#[inline(always)]
pub fn alloc_slice_fill_with<T>(&self, len: usize, fill: impl FnMut(usize) -> T) -> &mut [T] {
self.arena.alloc_slice_fill_with(len, fill)
}
}
#[derive(Clone, Copy, Debug)]
pub struct DynamicOneSeedBatch<'arena> {
pub base: DynamicOrder2<'arena>,
eps: &'arena [DynamicOrder2<'arena>],
}
impl<'arena> DynamicOneSeedBatch<'arena> {
#[inline(always)]
#[must_use]
pub fn seed_directions(
x: f64,
axis: usize,
dimension: usize,
workspace: &'arena DynamicJetBatchWorkspace,
mut direction_at: impl FnMut(usize) -> f64,
) -> Self {
let eps = workspace
.arena
.alloc_slice_fill_with(workspace.lanes, |lane| {
DynamicOrder2::constant(direction_at(lane), dimension, &workspace.arena)
});
Self {
base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
eps,
}
}
#[inline(always)]
#[must_use]
pub fn lanes(&self) -> usize {
self.eps.len()
}
#[inline(always)]
#[must_use]
pub fn contracted_third(&self, lane: usize) -> &[f64] {
self.eps[lane].h()
}
#[inline(always)]
fn assert_compatible(&self, other: &Self) {
self.base.assert_compatible(&other.base);
assert_eq!(
self.eps.len(),
other.eps.len(),
"dynamic one-seed batch lane mismatch"
);
}
}
impl<'arena> RuntimeJetScalar<'arena> for DynamicOneSeedBatch<'arena> {
type Workspace = DynamicJetBatchWorkspace;
#[inline(always)]
fn constant(c: f64, dimension: usize, workspace: &'arena DynamicJetBatchWorkspace) -> Self {
let eps = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
DynamicOrder2::constant(0.0, dimension, &workspace.arena)
});
Self {
base: DynamicOrder2::constant(c, dimension, &workspace.arena),
eps,
}
}
#[inline(always)]
fn variable(
x: f64,
axis: usize,
dimension: usize,
workspace: &'arena DynamicJetBatchWorkspace,
) -> Self {
let eps = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
DynamicOrder2::constant(0.0, dimension, &workspace.arena)
});
Self {
base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
eps,
}
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
let eps = self
.base
.arena
.alloc_slice_fill_with(self.lanes(), |_| self.base.constant_like(0.0));
Self {
base: self.base.constant_like(c),
eps,
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
base: self.base.with_value(value),
eps: self.eps,
}
}
#[inline(always)]
fn dimension(&self) -> usize {
self.base.dimension()
}
#[inline(always)]
fn value(&self) -> f64 {
self.base.value()
}
#[inline(always)]
fn add(&self, other: &Self) -> Self {
self.assert_compatible(other);
let eps = self
.base
.arena
.alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].add(&other.eps[lane]));
Self {
base: self.base.add(&other.base),
eps,
}
}
#[inline(always)]
fn sub(&self, other: &Self) -> Self {
self.assert_compatible(other);
let eps = self
.base
.arena
.alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].sub(&other.eps[lane]));
Self {
base: self.base.sub(&other.base),
eps,
}
}
#[inline(always)]
fn mul(&self, other: &Self) -> Self {
self.assert_compatible(other);
let eps = self
.base
.arena
.alloc_slice_fill_with(self.eps.len(), |lane| {
self.base
.mul(&other.eps[lane])
.add(&self.eps[lane].mul(&other.base))
});
Self {
base: self.base.mul(&other.base),
eps,
}
}
#[inline(always)]
fn neg(&self) -> Self {
self.scale(-1.0)
}
#[inline(always)]
fn scale(&self, scale: f64) -> Self {
let eps = self
.base
.arena
.alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].scale(scale));
Self {
base: self.base.scale(scale),
eps,
}
}
#[inline(always)]
fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
let fprime = self.base.compose_unary([
derivatives[1],
derivatives[2],
derivatives[3],
derivatives[4],
derivatives[4],
]);
let eps = self
.base
.arena
.alloc_slice_fill_with(self.eps.len(), |lane| fprime.mul(&self.eps[lane]));
Self {
base: self.base.compose_unary(derivatives),
eps,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct DynamicTwoSeedBatch<'arena> {
pub base: DynamicOrder2<'arena>,
eps: &'arena [DynamicOrder2<'arena>],
del: &'arena [DynamicOrder2<'arena>],
eps_del: &'arena [DynamicOrder2<'arena>],
}
impl<'arena> DynamicTwoSeedBatch<'arena> {
#[inline(always)]
#[must_use]
pub fn seed_direction_pairs(
x: f64,
axis: usize,
dimension: usize,
workspace: &'arena DynamicJetBatchWorkspace,
mut direction_pair_at: impl FnMut(usize) -> (f64, f64),
) -> Self {
let directions = workspace
.arena
.alloc_slice_fill_with(workspace.lanes, |lane| direction_pair_at(lane));
let eps = workspace
.arena
.alloc_slice_fill_with(workspace.lanes, |lane| {
DynamicOrder2::constant(directions[lane].0, dimension, &workspace.arena)
});
let del = workspace
.arena
.alloc_slice_fill_with(workspace.lanes, |lane| {
DynamicOrder2::constant(directions[lane].1, dimension, &workspace.arena)
});
let eps_del = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
DynamicOrder2::constant(0.0, dimension, &workspace.arena)
});
Self {
base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
eps,
del,
eps_del,
}
}
#[inline(always)]
#[must_use]
pub fn lanes(&self) -> usize {
self.eps.len()
}
#[inline(always)]
#[must_use]
pub fn contracted_fourth(&self, lane: usize) -> &[f64] {
self.eps_del[lane].h()
}
#[inline(always)]
fn assert_compatible(&self, other: &Self) {
self.base.assert_compatible(&other.base);
assert_eq!(
self.eps.len(),
other.eps.len(),
"dynamic two-seed batch lane mismatch"
);
assert_eq!(
self.del.len(),
self.eps.len(),
"dynamic two-seed batch delta mismatch"
);
assert_eq!(
self.eps_del.len(),
self.eps.len(),
"dynamic two-seed batch cross mismatch"
);
}
}
impl<'arena> RuntimeJetScalar<'arena> for DynamicTwoSeedBatch<'arena> {
type Workspace = DynamicJetBatchWorkspace;
#[inline(always)]
fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self {
let zero = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
DynamicOrder2::constant(0.0, dimension, &workspace.arena)
});
Self {
base: DynamicOrder2::constant(c, dimension, &workspace.arena),
eps: zero,
del: zero,
eps_del: zero,
}
}
#[inline(always)]
fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self {
let zero = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
DynamicOrder2::constant(0.0, dimension, &workspace.arena)
});
Self {
base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
eps: zero,
del: zero,
eps_del: zero,
}
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
let zero = self
.base
.arena
.alloc_slice_fill_with(self.lanes(), |_| self.base.constant_like(0.0));
Self {
base: self.base.constant_like(c),
eps: zero,
del: zero,
eps_del: zero,
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
base: self.base.with_value(value),
eps: self.eps,
del: self.del,
eps_del: self.eps_del,
}
}
#[inline(always)]
fn dimension(&self) -> usize {
self.base.dimension()
}
#[inline(always)]
fn value(&self) -> f64 {
self.base.value()
}
#[inline(always)]
fn add(&self, other: &Self) -> Self {
self.assert_compatible(other);
let arena = self.base.arena;
let eps =
arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].add(&other.eps[lane]));
let del =
arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].add(&other.del[lane]));
let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
self.eps_del[lane].add(&other.eps_del[lane])
});
Self {
base: self.base.add(&other.base),
eps,
del,
eps_del,
}
}
#[inline(always)]
fn sub(&self, other: &Self) -> Self {
self.assert_compatible(other);
let arena = self.base.arena;
let eps =
arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].sub(&other.eps[lane]));
let del =
arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].sub(&other.del[lane]));
let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
self.eps_del[lane].sub(&other.eps_del[lane])
});
Self {
base: self.base.sub(&other.base),
eps,
del,
eps_del,
}
}
#[inline(always)]
fn mul(&self, other: &Self) -> Self {
self.assert_compatible(other);
let arena = self.base.arena;
let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| {
self.base
.mul(&other.eps[lane])
.add(&self.eps[lane].mul(&other.base))
});
let del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
self.base
.mul(&other.del[lane])
.add(&self.del[lane].mul(&other.base))
});
let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
self.base
.mul(&other.eps_del[lane])
.add(&self.eps[lane].mul(&other.del[lane]))
.add(&self.del[lane].mul(&other.eps[lane]))
.add(&self.eps_del[lane].mul(&other.base))
});
Self {
base: self.base.mul(&other.base),
eps,
del,
eps_del,
}
}
#[inline(always)]
fn neg(&self) -> Self {
self.scale(-1.0)
}
#[inline(always)]
fn scale(&self, scale: f64) -> Self {
let arena = self.base.arena;
let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].scale(scale));
let del = arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].scale(scale));
let eps_del =
arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps_del[lane].scale(scale));
Self {
base: self.base.scale(scale),
eps,
del,
eps_del,
}
}
#[inline(always)]
fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
let arena = self.base.arena;
let fprime = self.base.compose_unary([
derivatives[1],
derivatives[2],
derivatives[3],
derivatives[4],
derivatives[4],
]);
let fsecond = self.base.compose_unary([
derivatives[2],
derivatives[3],
derivatives[4],
derivatives[4],
derivatives[4],
]);
let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| fprime.mul(&self.eps[lane]));
let del = arena.alloc_slice_fill_with(self.lanes(), |lane| fprime.mul(&self.del[lane]));
let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
fsecond
.mul(&self.eps[lane])
.mul(&self.del[lane])
.add(&fprime.mul(&self.eps_del[lane]))
});
Self {
base: self.base.compose_unary(derivatives),
eps,
del,
eps_del,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct DynamicTwoSeed<'arena> {
pub base: DynamicOrder2<'arena>,
pub eps: DynamicOrder2<'arena>,
pub del: DynamicOrder2<'arena>,
pub eps_del: DynamicOrder2<'arena>,
}
impl<'arena> DynamicTwoSeed<'arena> {
#[inline(always)]
#[must_use]
pub fn seed(
x: f64,
axis: usize,
u_axis: f64,
v_axis: f64,
dimension: usize,
arena: &'arena DynamicJetArena,
) -> Self {
Self {
base: DynamicOrder2::variable(x, axis, dimension, arena),
eps: DynamicOrder2::constant(u_axis, dimension, arena),
del: DynamicOrder2::constant(v_axis, dimension, arena),
eps_del: DynamicOrder2::constant(0.0, dimension, arena),
}
}
#[inline(always)]
#[must_use]
pub fn contracted_fourth(&self) -> &[f64] {
self.eps_del.h()
}
}
impl<'arena> RuntimeJetScalar<'arena> for DynamicTwoSeed<'arena> {
type Workspace = DynamicJetArena;
#[inline(always)]
fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
Self {
base: DynamicOrder2::constant(c, dimension, arena),
eps: DynamicOrder2::constant(0.0, dimension, arena),
del: DynamicOrder2::constant(0.0, dimension, arena),
eps_del: DynamicOrder2::constant(0.0, dimension, arena),
}
}
#[inline(always)]
fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
Self {
base: DynamicOrder2::variable(x, axis, dimension, arena),
eps: DynamicOrder2::constant(0.0, dimension, arena),
del: DynamicOrder2::constant(0.0, dimension, arena),
eps_del: DynamicOrder2::constant(0.0, dimension, arena),
}
}
#[inline(always)]
fn constant_like(&self, c: f64) -> Self {
Self {
base: self.base.constant_like(c),
eps: self.eps.constant_like(0.0),
del: self.del.constant_like(0.0),
eps_del: self.eps_del.constant_like(0.0),
}
}
#[inline(always)]
fn with_value(&self, value: f64) -> Self {
Self {
base: self.base.with_value(value),
eps: self.eps,
del: self.del,
eps_del: self.eps_del,
}
}
#[inline(always)]
fn dimension(&self) -> usize {
self.base.dimension()
}
#[inline(always)]
fn value(&self) -> f64 {
self.base.value()
}
#[inline(always)]
fn add(&self, o: &Self) -> Self {
Self {
base: self.base.add(&o.base),
eps: self.eps.add(&o.eps),
del: self.del.add(&o.del),
eps_del: self.eps_del.add(&o.eps_del),
}
}
#[inline(always)]
fn sub(&self, o: &Self) -> Self {
Self {
base: self.base.sub(&o.base),
eps: self.eps.sub(&o.eps),
del: self.del.sub(&o.del),
eps_del: self.eps_del.sub(&o.eps_del),
}
}
#[inline(always)]
fn mul(&self, o: &Self) -> Self {
let base = self.base.mul(&o.base);
let eps = self.base.mul(&o.eps).add(&self.eps.mul(&o.base));
let del = self.base.mul(&o.del).add(&self.del.mul(&o.base));
let eps_del = self
.base
.mul(&o.eps_del)
.add(&self.eps.mul(&o.del))
.add(&self.del.mul(&o.eps))
.add(&self.eps_del.mul(&o.base));
Self {
base,
eps,
del,
eps_del,
}
}
#[inline(always)]
fn neg(&self) -> Self {
Self {
base: self.base.neg(),
eps: self.eps.neg(),
del: self.del.neg(),
eps_del: self.eps_del.neg(),
}
}
#[inline(always)]
fn scale(&self, s: f64) -> Self {
Self {
base: self.base.scale(s),
eps: self.eps.scale(s),
del: self.del.scale(s),
eps_del: self.eps_del.scale(s),
}
}
#[inline(always)]
fn compose_unary(&self, d: [f64; 5]) -> Self {
let base = self.base.compose_unary(d);
let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]);
let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]);
let eps = fprime.mul(&self.eps);
let del = fprime.mul(&self.del);
let eps_del = fsecond
.mul(&self.eps)
.mul(&self.del)
.add(&fprime.mul(&self.eps_del));
Self {
base,
eps,
del,
eps_del,
}
}
}
impl<const K: usize> std::ops::Add for Order2<K> {
type Output = Self;
#[inline]
fn add(self, o: Self) -> Self {
Order2(self.0 + o.0)
}
}
impl<const K: usize> std::ops::Add<f64> for Order2<K> {
type Output = Self;
#[inline]
fn add(self, c: f64) -> Self {
Order2(self.0 + c)
}
}
impl<const K: usize> std::ops::Sub for Order2<K> {
type Output = Self;
#[inline]
fn sub(self, o: Self) -> Self {
Order2(self.0 + o.0.scale(-1.0))
}
}
impl<const K: usize> std::ops::Sub<f64> for Order2<K> {
type Output = Self;
#[inline]
fn sub(self, c: f64) -> Self {
Order2(self.0 + (-c))
}
}
impl<const K: usize> std::ops::Mul for Order2<K> {
type Output = Self;
#[inline]
fn mul(self, o: Self) -> Self {
Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
}
}
impl<const K: usize> std::ops::Mul<f64> for Order2<K> {
type Output = Self;
#[inline]
fn mul(self, c: f64) -> Self {
Order2(self.0.scale(c))
}
}
impl<const K: usize> std::ops::Neg for Order2<K> {
type Output = Self;
#[inline]
fn neg(self) -> Self {
Order2(self.0.scale(-1.0))
}
}
pub fn filtered_implicit_solve_scalar<const K: usize, S: JetScalar<K>>(
a0: f64,
inv_fa: f64,
iters: usize,
f: impl Fn(&S) -> S,
) -> S {
let mut a = S::constant(a0);
for _ in 0..iters {
let residual = f(&a);
a = a.sub(&residual.scale(inv_fa));
}
a
}
pub fn filtered_implicit_solve_runtime_scalar<'arena, S: RuntimeJetScalar<'arena>>(
a0: f64,
inv_fa: f64,
iters: usize,
dimension: usize,
workspace: &'arena S::Workspace,
f: impl Fn(&S) -> S,
) -> S {
let mut a = S::constant(a0, dimension, workspace);
for _ in 0..iters {
let residual = f(&a);
a = a.sub(&residual.scale(inv_fa));
}
a
}
pub trait HessianPattern<const K: usize, const H: usize> {
const PAIRS: [(usize, usize); H];
const PAIR_BITS: [[u128; K]; K];
}
pub const fn hessian_pair_bits<const K: usize, const H: usize>(
pairs: [(usize, usize); H],
) -> [[u128; K]; K] {
let mut table = [[0u128; K]; K];
let mut slot = 0;
while slot < H {
let (i, j) = pairs[slot];
let bit = 1u128 << slot;
table[i][j] = bit;
table[j][i] = bit;
slot += 1;
}
table
}
#[derive(Debug)]
pub struct PatternedOrder2<P, const K: usize, const H: usize> {
v: f64,
g: [f64; K],
h: [f64; H],
gradient_mask: u128,
hessian_mask: u128,
pattern: std::marker::PhantomData<fn() -> P>,
}
impl<P, const K: usize, const H: usize> Copy for PatternedOrder2<P, K, H> {}
impl<P, const K: usize, const H: usize> Clone for PatternedOrder2<P, K, H> {
fn clone(&self) -> Self {
*self
}
}
impl<P, const K: usize, const H: usize> PatternedOrder2<P, K, H>
where
P: HessianPattern<K, H>,
{
#[inline]
#[must_use]
pub fn g(&self) -> [f64; K] {
self.g
}
#[inline]
#[must_use]
pub fn h(&self) -> [[f64; K]; K] {
let mut dense = [[0.0; K]; K];
for (slot, &(i, j)) in P::PAIRS.iter().enumerate() {
dense[i][j] = self.h[slot];
dense[j][i] = self.h[slot];
}
dense
}
#[inline]
fn pair_mask_between(left: u128, right: u128) -> u128 {
let mut result = 0u128;
let mut left_axes = left;
while left_axes != 0 {
let i = left_axes.trailing_zeros() as usize;
left_axes &= left_axes - 1;
let mut right_axes = right;
while right_axes != 0 {
let j = right_axes.trailing_zeros() as usize;
right_axes &= right_axes - 1;
result |= P::PAIR_BITS[i][j];
}
}
result
}
}
impl<P, const K: usize, const H: usize> JetScalar<K> for PatternedOrder2<P, K, H>
where
P: HessianPattern<K, H>,
{
#[inline]
fn constant(c: f64) -> Self {
Self {
v: c,
g: [0.0; K],
h: [0.0; H],
gradient_mask: 0,
hessian_mask: 0,
pattern: std::marker::PhantomData,
}
}
#[inline]
fn variable(x: f64, axis: usize) -> Self {
let mut out = Self::constant(x);
if axis < K {
out.g[axis] = 1.0;
out.gradient_mask = 1u128 << axis;
}
out
}
}
impl<P, const K: usize, const H: usize> crate::nested_dual::JetField for PatternedOrder2<P, K, H>
where
P: HessianPattern<K, H>,
{
#[inline]
fn value(&self) -> f64 {
self.v
}
#[inline]
fn add(&self, other: &Self) -> Self {
let mut out = Self::constant(self.v + other.v);
out.gradient_mask = self.gradient_mask | other.gradient_mask;
let mut gradient_mask = out.gradient_mask;
while gradient_mask != 0 {
let i = gradient_mask.trailing_zeros() as usize;
gradient_mask &= gradient_mask - 1;
out.g[i] = self.g[i] + other.g[i];
}
out.hessian_mask = self.hessian_mask | other.hessian_mask;
let mut hessian_mask = out.hessian_mask;
while hessian_mask != 0 {
let slot = hessian_mask.trailing_zeros() as usize;
hessian_mask &= hessian_mask - 1;
out.h[slot] = self.h[slot] + other.h[slot];
}
out
}
#[inline]
fn sub(&self, other: &Self) -> Self {
let mut out = Self::constant(self.v - other.v);
out.gradient_mask = self.gradient_mask | other.gradient_mask;
let mut gradient_mask = out.gradient_mask;
while gradient_mask != 0 {
let i = gradient_mask.trailing_zeros() as usize;
gradient_mask &= gradient_mask - 1;
out.g[i] = self.g[i] - other.g[i];
}
out.hessian_mask = self.hessian_mask | other.hessian_mask;
let mut hessian_mask = out.hessian_mask;
while hessian_mask != 0 {
let slot = hessian_mask.trailing_zeros() as usize;
hessian_mask &= hessian_mask - 1;
out.h[slot] = self.h[slot] - other.h[slot];
}
out
}
#[inline]
fn mul(&self, other: &Self) -> Self {
let mut out = Self::constant(self.v * other.v);
out.gradient_mask = self.gradient_mask | other.gradient_mask;
let mut gradient_mask = out.gradient_mask;
while gradient_mask != 0 {
let i = gradient_mask.trailing_zeros() as usize;
gradient_mask &= gradient_mask - 1;
out.g[i] = self.v * other.g[i] + self.g[i] * other.v;
}
out.hessian_mask = self.hessian_mask
| other.hessian_mask
| Self::pair_mask_between(self.gradient_mask, other.gradient_mask);
let mut hessian_mask = out.hessian_mask;
while hessian_mask != 0 {
let slot = hessian_mask.trailing_zeros() as usize;
hessian_mask &= hessian_mask - 1;
let (i, j) = P::PAIRS[slot];
out.h[slot] = self.v * other.h[slot]
+ self.g[i] * other.g[j]
+ self.g[j] * other.g[i]
+ self.h[slot] * other.v;
}
out
}
#[inline]
fn neg(&self) -> Self {
self.scale(-1.0)
}
#[inline]
fn scale(&self, scale: f64) -> Self {
let mut out = Self::constant(self.v * scale);
out.gradient_mask = self.gradient_mask;
let mut gradient_mask = out.gradient_mask;
while gradient_mask != 0 {
let i = gradient_mask.trailing_zeros() as usize;
gradient_mask &= gradient_mask - 1;
out.g[i] = self.g[i] * scale;
}
out.hessian_mask = self.hessian_mask;
let mut hessian_mask = out.hessian_mask;
while hessian_mask != 0 {
let slot = hessian_mask.trailing_zeros() as usize;
hessian_mask &= hessian_mask - 1;
out.h[slot] = self.h[slot] * scale;
}
out
}
#[inline]
fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
let mut out = Self::constant(derivatives[0]);
out.gradient_mask = self.gradient_mask;
let mut gradient_mask = out.gradient_mask;
while gradient_mask != 0 {
let i = gradient_mask.trailing_zeros() as usize;
gradient_mask &= gradient_mask - 1;
out.g[i] = derivatives[1] * self.g[i];
}
out.hessian_mask =
self.hessian_mask | Self::pair_mask_between(self.gradient_mask, self.gradient_mask);
let mut hessian_mask = out.hessian_mask;
while hessian_mask != 0 {
let slot = hessian_mask.trailing_zeros() as usize;
hessian_mask &= hessian_mask - 1;
let (i, j) = P::PAIRS[slot];
out.h[slot] = derivatives[2] * self.g[i] * self.g[j] + derivatives[1] * self.h[slot];
}
out
}
}
#[derive(Clone, Copy, Debug)]
pub struct Order2<const K: usize>(pub crate::jet_tower::Tower2<K>);
impl<const K: usize> Order2<K> {
#[inline]
#[must_use]
pub fn g(&self) -> &[f64; K] {
&self.0.g
}
#[inline]
#[must_use]
pub fn h(&self) -> &[[f64; K]; K] {
&self.0.h
}
#[inline]
#[must_use]
pub fn into_channels(self) -> (f64, [f64; K], [[f64; K]; K]) {
let crate::jet_tower::Tower2 { v, g, h } = self.0;
(v, g, h)
}
}
impl<const K: usize> JetScalar<K> for Order2<K> {
fn constant(c: f64) -> Self {
Order2(crate::jet_tower::Tower2::constant(c))
}
fn variable(x: f64, axis: usize) -> Self {
Order2(crate::jet_tower::Tower2::variable(x, axis))
}
#[inline(always)]
fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
inputs: &[Self],
coefficients: &C,
) -> Self {
assert_eq!(inputs.len(), coefficients.dimension());
let input_dimension = inputs.len();
assert!(input_dimension <= K);
let mut values = [0.0; K];
for axis in 0..input_dimension {
values[axis] = inputs[axis].0.v;
}
let mut projected = [0.0; K];
coefficients.multiply(
&values[..input_dimension],
&mut projected[..input_dimension],
);
let mut out = crate::jet_tower::Tower2::zero();
for axis in 0..input_dimension {
out.v += values[axis] * projected[axis];
}
for primary in 0..K {
let mut channel = 0.0;
for axis in 0..input_dimension {
channel += projected[axis] * inputs[axis].0.g[primary];
}
out.g[primary] = 2.0 * channel;
}
let mut input_gradient = [0.0; K];
let mut projected_gradient = [0.0; K];
for primary_b in 0..K {
for row in 0..input_dimension {
input_gradient[row] = inputs[row].0.g[primary_b];
}
coefficients.multiply(
&input_gradient[..input_dimension],
&mut projected_gradient[..input_dimension],
);
for primary_a in 0..=primary_b {
let mut inherited = 0.0;
let mut curvature = 0.0;
for row in 0..input_dimension {
inherited += projected[row] * inputs[row].0.h[primary_a][primary_b];
curvature += inputs[row].0.g[primary_a] * projected_gradient[row];
}
let channel = 2.0 * (inherited + curvature);
out.h[primary_a][primary_b] = channel;
out.h[primary_b][primary_a] = channel;
}
}
Order2(out)
}
#[inline(always)]
fn linear_combination(inputs: &[Self], weights: &[f64]) -> Self {
assert_eq!(inputs.len(), weights.len());
let mut out = crate::jet_tower::Tower2::zero();
for (input, &weight) in inputs.iter().zip(weights) {
out.v += input.0.v * weight;
}
for primary in 0..K {
for (input, &weight) in inputs.iter().zip(weights) {
out.g[primary] += input.0.g[primary] * weight;
}
for other in primary..K {
for (input, &weight) in inputs.iter().zip(weights) {
out.h[primary][other] += input.0.h[primary][other] * weight;
}
out.h[other][primary] = out.h[primary][other];
}
}
Order2(out)
}
#[inline(always)]
fn add_constant(&self, constant: f64) -> Self {
let mut out = *self;
out.0.v += constant;
out
}
#[inline(always)]
fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
let mut out = crate::jet_tower::Tower2::zero();
out.v = self.0.v * right.0.v + addend.0.v;
for primary in 0..K {
out.g[primary] =
self.0.v * right.0.g[primary] + self.0.g[primary] * right.0.v + addend.0.g[primary];
for other in primary..K {
let channel = self.0.v * right.0.h[primary][other]
+ self.0.g[primary] * right.0.g[other]
+ self.0.g[other] * right.0.g[primary]
+ self.0.h[primary][other] * right.0.v
+ addend.0.h[primary][other];
out.h[primary][other] = channel;
out.h[other][primary] = channel;
}
}
Order2(out)
}
#[inline(always)]
fn product(&self, right: &Self) -> Self {
let mut out = crate::jet_tower::Tower2::zero();
out.v = self.0.v * right.0.v;
for primary in 0..K {
out.g[primary] = self.0.v * right.0.g[primary] + self.0.g[primary] * right.0.v;
for other in primary..K {
let channel = self.0.v * right.0.h[primary][other]
+ self.0.g[primary] * right.0.g[other]
+ self.0.g[other] * right.0.g[primary]
+ self.0.h[primary][other] * right.0.v;
out.h[primary][other] = channel;
out.h[other][primary] = channel;
}
}
Order2(out)
}
#[inline(always)]
fn affine_compose(
&self,
input_scale: f64,
input_shift: f64,
derivative_stack: [f64; 5],
) -> Self {
assert!(input_shift.is_finite(), "affine input shift must be finite");
let first = derivative_stack[1] * input_scale;
let second = derivative_stack[2] * input_scale * input_scale;
let mut out = crate::jet_tower::Tower2::zero();
out.v = derivative_stack[0];
for primary in 0..K {
out.g[primary] = first * self.0.g[primary];
for other in primary..K {
let channel =
first * self.0.h[primary][other] + second * self.0.g[primary] * self.0.g[other];
out.h[primary][other] = channel;
out.h[other][primary] = channel;
}
}
Order2(out)
}
#[inline(always)]
fn affine_composed_sum(
inputs: &[Self],
input_scales: &[f64],
derivative_stacks: &[[f64; 5]],
) -> Self {
assert_eq!(inputs.len(), input_scales.len());
assert_eq!(inputs.len(), derivative_stacks.len());
let mut out = crate::jet_tower::Tower2::zero();
for ((input, &input_scale), stack) in inputs.iter().zip(input_scales).zip(derivative_stacks)
{
let first = stack[1] * input_scale;
let second = stack[2] * input_scale * input_scale;
out.v += stack[0];
for primary in 0..K {
out.g[primary] += first * input.0.g[primary];
for other in primary..K {
out.h[primary][other] += first * input.0.h[primary][other]
+ second * input.0.g[primary] * input.0.g[other];
}
}
}
for primary in 0..K {
for other in primary + 1..K {
out.h[other][primary] = out.h[primary][other];
}
}
Order2(out)
}
#[inline(always)]
fn shared_multiply_add_affine_composed_sum<const N: usize>(
lefts: &[&Self; N],
right: &Self,
addend: &Self,
addend_scales: &[f64; N],
input_scales: &[f64; N],
derivative_stacks: &[[f64; 5]; N],
) -> Self {
let (representatives, term_sources, source_count) =
canonical_shared_source_schedule::<N>(|term, representative| {
std::ptr::eq(lefts[term], lefts[representative])
&& addend_scales[term] == addend_scales[representative]
});
let (value, source_derivatives) =
aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
let mut source_gradients = [[0.0; K]; N];
let mut out = crate::jet_tower::Tower2::zero();
out.v = value;
let mut right_first = 0.0;
let mut addend_first = 0.0;
for source in 0..source_count {
let term = representatives[source];
let first = source_derivatives[source][1];
right_first += first * lefts[term].0.v;
addend_first += first * addend_scales[term];
for primary in 0..K {
let product_gradient =
lefts[term].0.v * right.0.g[primary] + lefts[term].0.g[primary] * right.0.v;
let inner_gradient = if addend_scales[term] == 0.0 {
product_gradient
} else if addend_scales[term] == 1.0 {
product_gradient + addend.0.g[primary]
} else {
product_gradient + addend_scales[term] * addend.0.g[primary]
};
source_gradients[source][primary] = inner_gradient;
out.g[primary] += first * lefts[term].0.g[primary] * right.0.v;
}
}
if N != 0 {
for primary in 0..K {
out.g[primary] += right_first * right.0.g[primary];
}
}
let addend_live = addend_scales.iter().any(|&scale| scale != 0.0);
if addend_live {
for primary in 0..K {
out.g[primary] += addend_first * addend.0.g[primary];
}
}
for primary in 0..K {
for other in primary..K {
let mut channel = if N == 0 {
0.0
} else {
right_first * right.0.h[primary][other]
};
if addend_live {
channel += addend_first * addend.0.h[primary][other];
}
for source in 0..source_count {
let term = representatives[source];
let local_product_hessian = lefts[term].0.g[primary] * right.0.g[other]
+ lefts[term].0.g[other] * right.0.g[primary]
+ lefts[term].0.h[primary][other] * right.0.v;
channel += source_derivatives[source][1] * local_product_hessian
+ source_derivatives[source][2]
* source_gradients[source][primary]
* source_gradients[source][other];
}
out.h[primary][other] = channel;
out.h[other][primary] = channel;
}
}
Order2(out)
}
#[inline(always)]
fn composed_sum(inputs: &[Self], derivative_stacks: &[[f64; 5]]) -> Self {
assert_eq!(inputs.len(), derivative_stacks.len());
let mut out = crate::jet_tower::Tower2::zero();
for (input, stack) in inputs.iter().zip(derivative_stacks) {
out.v += stack[0];
for primary in 0..K {
out.g[primary] += stack[1] * input.0.g[primary];
for other in primary..K {
out.h[primary][other] += stack[1] * input.0.h[primary][other]
+ stack[2] * input.0.g[primary] * input.0.g[other];
}
}
}
for primary in 0..K {
for other in primary + 1..K {
out.h[other][primary] = out.h[primary][other];
}
}
Order2(out)
}
}
impl<const K: usize> crate::nested_dual::JetField for Order2<K> {
fn value(&self) -> f64 {
self.0.v
}
fn add(&self, o: &Self) -> Self {
Order2(self.0 + o.0)
}
fn sub(&self, o: &Self) -> Self {
Order2(self.0 + o.0.scale(-1.0))
}
fn mul(&self, o: &Self) -> Self {
Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
}
fn neg(&self) -> Self {
Order2(self.0.scale(-1.0))
}
fn scale(&self, s: f64) -> Self {
Order2(self.0.scale(s))
}
fn compose_unary(&self, d: [f64; 5]) -> Self {
Order2(self.0.compose_unary([d[0], d[1], d[2]]))
}
fn constant_like(&self, v: f64) -> Self {
<Self as JetScalar<K>>::constant(v)
}
fn with_value(&self, v: f64) -> Self {
let mut out = *self;
out.0.v = v;
out
}
}
#[derive(Clone, Copy, Debug)]
pub struct MappedOrder2Accumulator<const K: usize> {
value: f64,
gradient: [f64; K],
hessian: [[f64; K]; K],
}
#[derive(Clone, Copy, Debug)]
pub struct StaticOrder2Atom<
const N: usize,
const H: usize,
const GRADIENT_BITS: u128,
const HESSIAN_BITS: u128,
> {
value: f64,
gradient: [f64; N],
hessian: [f64; H],
}
impl<const N: usize, const H: usize, const G: u128, const Q: u128> StaticOrder2Atom<N, H, G, Q> {
#[inline(always)]
#[must_use]
pub fn new(value: f64, gradient: [f64; N], hessian: [f64; H]) -> Self {
assert!(H == N * (N + 1) / 2, "invalid packed order-two shape");
assert!(N <= 128 && H <= 128, "static atom sparsity mask overflow");
Self {
value,
gradient,
hessian,
}
}
#[inline(always)]
#[must_use]
pub fn value(&self) -> f64 {
self.value
}
#[inline(always)]
#[must_use]
pub fn gradient(&self) -> [f64; N] {
self.gradient
}
#[inline(always)]
#[must_use]
pub fn hessian_at(&self, row: usize, column: usize) -> f64 {
assert!(
row < N && column < N,
"static atom Hessian axis out of range"
);
let (row, column) = if row <= column {
(row, column)
} else {
(column, row)
};
let index = row * (2 * N - row + 1) / 2 + column - row;
self.hessian[index]
}
}
pub trait Order2AtomChannels<const N: usize> {
const GRADIENT_BITS: u128;
const HESSIAN_BITS: u128;
fn gradient_at(&self, axis: usize) -> f64;
fn hessian_at(&self, row: usize, column: usize) -> f64;
}
impl<const N: usize> Order2AtomChannels<N> for Order2<N> {
const GRADIENT_BITS: u128 = low_mask(N);
const HESSIAN_BITS: u128 = low_mask(N * (N + 1) / 2);
#[inline(always)]
fn gradient_at(&self, axis: usize) -> f64 {
self.0.g[axis]
}
#[inline(always)]
fn hessian_at(&self, row: usize, column: usize) -> f64 {
self.0.h[row][column]
}
}
impl<const N: usize, const H: usize, const G: u128, const Q: u128> Order2AtomChannels<N>
for StaticOrder2Atom<N, H, G, Q>
{
const GRADIENT_BITS: u128 = G;
const HESSIAN_BITS: u128 = Q;
#[inline(always)]
fn gradient_at(&self, axis: usize) -> f64 {
self.gradient[axis]
}
#[inline(always)]
fn hessian_at(&self, row: usize, column: usize) -> f64 {
StaticOrder2Atom::hessian_at(self, row, column)
}
}
const fn low_mask(channels: usize) -> u128 {
if channels >= 128 {
u128::MAX
} else {
(1u128 << channels) - 1
}
}
impl<const K: usize> MappedOrder2Accumulator<K> {
#[inline(always)]
#[must_use]
pub fn zero() -> Self {
Self {
value: 0.0,
gradient: [0.0; K],
hessian: [[0.0; K]; K],
}
}
#[inline(always)]
pub fn add_composed<const N: usize, const H: usize, A: Order2AtomChannels<N>>(
&mut self,
atom: &A,
axes: [usize; N],
derivatives: [f64; 3],
value_add: bool,
gradient_add: [bool; N],
hessian_add: [bool; H],
) {
assert!(H == N * (N + 1) / 2, "invalid mapped Hessian write shape");
assert!(N <= 128 && H <= 128, "mapped atom sparsity mask overflow");
assert!(
axes.iter().all(|&axis| axis < K),
"mapped atom axis must be within the global primary dimension"
);
assert!(
axes.iter()
.enumerate()
.all(|(i, axis)| !axes[..i].contains(axis)),
"mapped atom axes must be injective"
);
if value_add {
self.value += derivatives[0];
} else {
self.value = derivatives[0];
}
let mut packed = 0;
for local_i in 0..N {
let global_i = axes[local_i];
if A::GRADIENT_BITS & (1u128 << local_i) != 0 {
let channel = derivatives[1] * atom.gradient_at(local_i);
if gradient_add[local_i] {
self.gradient[global_i] += channel;
} else {
self.gradient[global_i] = channel;
}
}
for local_j in local_i..N {
let global_j = axes[local_j];
let inner_live = A::HESSIAN_BITS & (1u128 << packed) != 0;
let outer_live = A::GRADIENT_BITS & (1u128 << local_i) != 0
&& A::GRADIENT_BITS & (1u128 << local_j) != 0;
let channel = if inner_live {
let inner = derivatives[1] * atom.hessian_at(local_i, local_j);
if outer_live {
inner
+ derivatives[2] * atom.gradient_at(local_i) * atom.gradient_at(local_j)
} else {
inner
}
} else if outer_live {
derivatives[2] * atom.gradient_at(local_i) * atom.gradient_at(local_j)
} else {
packed += 1;
continue;
};
if hessian_add[packed] {
self.hessian[global_i][global_j] += channel;
if global_i != global_j {
self.hessian[global_j][global_i] += channel;
}
} else {
self.hessian[global_i][global_j] = channel;
if global_i != global_j {
self.hessian[global_j][global_i] = channel;
}
}
packed += 1;
}
}
}
#[inline(always)]
#[must_use]
pub fn into_channels(self) -> (f64, [f64; K], [[f64; K]; K]) {
(self.value, self.gradient, self.hessian)
}
}
pub trait DynamicOrder2Term {
fn outer_first(&self) -> f64;
fn outer_second(&self) -> f64;
fn inner_gradient(&self, axis: usize) -> f64;
fn inner_hessian(&self, row: usize, column: usize) -> f64;
}
#[derive(Debug)]
pub struct DynamicOrder2Accumulator {
value: f64,
gradient: Vec<f64>,
hessian: Vec<f64>,
}
impl DynamicOrder2Accumulator {
#[inline(always)]
#[must_use]
pub fn from_composed_sum<T: DynamicOrder2Term, const N: usize>(
dimension: usize,
value: f64,
terms: &[T; N],
) -> Self {
let mut gradient = vec![0.0; dimension];
let mut hessian = vec![0.0; dimension * dimension];
for axis in 0..dimension {
let mut channel = 0.0;
for term in terms {
channel += term.outer_first() * term.inner_gradient(axis);
}
gradient[axis] = channel;
}
for row in 0..dimension {
for column in row..dimension {
let mut channel = 0.0;
for term in terms {
let row_gradient = term.inner_gradient(row);
let column_gradient = term.inner_gradient(column);
channel += term.outer_second() * row_gradient * column_gradient
+ term.outer_first() * term.inner_hessian(row, column);
}
hessian[row * dimension + column] = channel;
hessian[column * dimension + row] = channel;
}
}
Self {
value,
gradient,
hessian,
}
}
#[inline(always)]
#[must_use]
pub fn into_channels(self) -> (f64, Vec<f64>, Vec<f64>) {
(self.value, self.gradient, self.hessian)
}
}
pub trait Lane: Copy {
const LANES: usize;
fn splat(x: f64) -> Self;
fn add(self, o: Self) -> Self;
fn sub(self, o: Self) -> Self;
fn mul(self, o: Self) -> Self;
fn lane(self, i: usize) -> f64;
fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3];
fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5];
}
impl Lane for f64 {
const LANES: usize = 1;
#[inline]
fn splat(x: f64) -> Self {
x
}
#[inline]
fn add(self, o: Self) -> Self {
self + o
}
#[inline]
fn sub(self, o: Self) -> Self {
self - o
}
#[inline]
fn mul(self, o: Self) -> Self {
self * o
}
#[inline]
fn lane(self, i: usize) -> f64 {
assert!(
i < <Self as Lane>::LANES,
"the f64 Lane carries one row; lane {i} does not exist"
);
self
}
#[inline]
fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
stack(self)
}
#[inline]
fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
stack(self)
}
}
impl Lane for wide::f64x4 {
const LANES: usize = 4;
#[inline]
fn splat(x: f64) -> Self {
wide::f64x4::splat(x)
}
#[inline]
fn add(self, o: Self) -> Self {
self + o
}
#[inline]
fn sub(self, o: Self) -> Self {
self - o
}
#[inline]
fn mul(self, o: Self) -> Self {
self * o
}
#[inline]
fn lane(self, i: usize) -> f64 {
self.to_array()[i]
}
#[inline]
fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
let a = self.to_array();
let mut d0 = [0.0_f64; 4];
let mut d1 = [0.0_f64; 4];
let mut d2 = [0.0_f64; 4];
for i in 0..4 {
let s = stack(a[i]);
d0[i] = s[0];
d1[i] = s[1];
d2[i] = s[2];
}
[
wide::f64x4::new(d0),
wide::f64x4::new(d1),
wide::f64x4::new(d2),
]
}
#[inline]
fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
let a = self.to_array();
let mut d = [[0.0_f64; 4]; 5];
for i in 0..4 {
let s = stack(a[i]);
for (k, dk) in d.iter_mut().enumerate() {
dk[i] = s[k];
}
}
[
wide::f64x4::new(d[0]),
wide::f64x4::new(d[1]),
wide::f64x4::new(d[2]),
wide::f64x4::new(d[3]),
wide::f64x4::new(d[4]),
]
}
}
#[derive(Clone, Copy, Debug)]
pub struct Order2Lane<L: Lane, const K: usize> {
pub v: L,
pub g: [L; K],
pub h: [[L; K]; K],
}
pub type Order2Batch<const K: usize> = Order2Lane<wide::f64x4, K>;
impl<L: Lane, const K: usize> Order2Lane<L, K> {
#[inline]
pub fn constant(c: L) -> Self {
Order2Lane {
v: c,
g: [L::splat(0.0); K],
h: [[L::splat(0.0); K]; K],
}
}
#[inline]
pub fn variable(value: L, axis: usize) -> Self {
let mut out = Self::constant(value);
out.g[axis] = L::splat(1.0);
out
}
#[inline]
pub fn add(&self, o: &Self) -> Self {
let mut out = *self;
out.v = self.v.add(o.v);
for i in 0..K {
out.g[i] = self.g[i].add(o.g[i]);
for j in 0..K {
out.h[i][j] = self.h[i][j].add(o.h[i][j]);
}
}
out
}
#[inline]
pub fn scale(&self, s: f64) -> Self {
let sl = L::splat(s);
let mut out = *self;
out.v = self.v.mul(sl);
for i in 0..K {
out.g[i] = self.g[i].mul(sl);
for j in 0..K {
out.h[i][j] = self.h[i][j].mul(sl);
}
}
out
}
#[inline]
pub fn sub(&self, o: &Self) -> Self {
self.add(&o.scale(-1.0))
}
#[inline]
pub fn neg(&self) -> Self {
self.scale(-1.0)
}
#[inline]
pub fn mul(&self, o: &Self) -> Self {
let a = self;
let b = o;
let mut out = Self::constant(a.v.mul(b.v));
for i in 0..K {
out.g[i] = a.v.mul(b.g[i]).add(a.g[i].mul(b.v));
}
for i in 0..K {
for j in i..K {
let hij =
a.v.mul(b.h[i][j])
.add(a.g[i].mul(b.g[j]))
.add(a.g[j].mul(b.g[i]))
.add(a.h[i][j].mul(b.v));
out.h[i][j] = hij;
out.h[j][i] = hij;
}
}
out
}
#[inline]
pub fn compose_unary(&self, d: [L; 3]) -> Self {
let mut out = Self::constant(d[0]);
for i in 0..K {
let mut acc = L::splat(0.0);
acc = acc.add(d[1].mul(self.g[i]));
out.g[i] = acc;
}
for i in 0..K {
for j in 0..K {
let mut acc = L::splat(0.0);
acc = acc.add(d[1].mul(self.h[i][j]));
acc = acc.add(d[2].mul(self.g[i]).mul(self.g[j]));
out.h[i][j] = acc;
}
}
out
}
#[inline]
pub fn exp(&self) -> Self {
let d = self.v.unary3(|u| {
let e = u.exp();
[e, e, e]
});
self.compose_unary(d)
}
#[inline]
pub fn ln(&self) -> Self {
let d = self.v.unary3(|u| {
let r = 1.0 / u;
[u.ln(), r, -r * r]
});
self.compose_unary(d)
}
#[inline]
pub fn sqrt(&self) -> Self {
let d = self.v.unary3(|u| {
let s = u.sqrt();
[s, 0.5 / s, -0.25 / (u * s)]
});
self.compose_unary(d)
}
#[inline]
pub fn recip(&self) -> Self {
let d = self.v.unary3(|u| {
let r = 1.0 / u;
let r2 = r * r;
[r, -r2, 2.0 * r2 * r]
});
self.compose_unary(d)
}
#[inline]
pub fn powf(&self, a: f64) -> Self {
let d = self.v.unary3(|u| {
[
u.powf(a),
a * u.powf(a - 1.0),
a * (a - 1.0) * u.powf(a - 2.0),
]
});
self.compose_unary(d)
}
}
impl<const K: usize> Order2Batch<K> {
#[inline]
#[must_use]
pub fn lane(&self, i: usize) -> Order2<K> {
let mut t = crate::jet_tower::Tower2::<K>::constant(self.v.lane(i));
for a in 0..K {
t.g[a] = self.g[a].lane(i);
for b in 0..K {
t.h[a][b] = self.h[a][b].lane(i);
}
}
Order2(t)
}
}
#[derive(Clone, Copy, Debug)]
pub struct Order1<const K: usize> {
pub v: f64,
pub g: [f64; K],
}
impl<const K: usize> Order1<K> {
#[inline]
#[must_use]
pub fn g(&self) -> &[f64; K] {
&self.g
}
#[inline]
#[must_use]
pub fn into_channels(self) -> (f64, [f64; K]) {
(self.v, self.g)
}
}
impl<const K: usize> JetScalar<K> for Order1<K> {
fn constant(c: f64) -> Self {
Order1 { v: c, g: [0.0; K] }
}
fn variable(x: f64, axis: usize) -> Self {
let mut g = [0.0; K];
g[axis] = 1.0;
Order1 { v: x, g }
}
}
impl<const K: usize> crate::nested_dual::JetField for Order1<K> {
fn value(&self) -> f64 {
self.v
}
fn add(&self, o: &Self) -> Self {
let mut g = self.g;
for i in 0..K {
g[i] += o.g[i];
}
Order1 { v: self.v + o.v, g }
}
fn sub(&self, o: &Self) -> Self {
self.add(&o.scale(-1.0))
}
fn mul(&self, o: &Self) -> Self {
let a = self;
let b = o;
let mut g = [0.0; K];
for i in 0..K {
g[i] = a.v * b.g[i] + a.g[i] * b.v;
}
Order1 { v: a.v * b.v, g }
}
fn neg(&self) -> Self {
self.scale(-1.0)
}
fn scale(&self, s: f64) -> Self {
let mut g = self.g;
for i in 0..K {
g[i] *= s;
}
Order1 { v: self.v * s, g }
}
fn compose_unary(&self, d: [f64; 5]) -> Self {
let mut g = [0.0; K];
for i in 0..K {
g[i] = d[1] * self.g[i];
}
Order1 { v: d[0], g }
}
}
#[derive(Clone, Copy, Debug)]
pub struct OneSeed<const K: usize> {
pub base: Order2<K>,
pub eps: Order2<K>,
}
impl<const K: usize> OneSeed<K> {
pub fn seed_direction(x: f64, axis: usize, u_axis: f64) -> Self {
OneSeed {
base: Order2::variable(x, axis),
eps: Order2::constant(u_axis),
}
}
pub fn contracted_third(&self) -> [[f64; K]; K] {
*self.eps.h()
}
}
impl<const K: usize> JetScalar<K> for OneSeed<K> {
fn constant(c: f64) -> Self {
OneSeed {
base: Order2::constant(c),
eps: Order2::constant(0.0),
}
}
fn variable(x: f64, axis: usize) -> Self {
OneSeed {
base: Order2::variable(x, axis),
eps: Order2::constant(0.0),
}
}
}
impl<const K: usize> crate::nested_dual::JetField for OneSeed<K> {
fn value(&self) -> f64 {
self.base.value()
}
fn add(&self, o: &Self) -> Self {
OneSeed {
base: self.base.add(&o.base),
eps: self.eps.add(&o.eps),
}
}
fn sub(&self, o: &Self) -> Self {
OneSeed {
base: self.base.sub(&o.base),
eps: self.eps.sub(&o.eps),
}
}
fn mul(&self, o: &Self) -> Self {
let ab = &self.base.0;
let ae = &self.eps.0;
let bb = &o.base.0;
let be = &o.eps.0;
let mut eps = crate::jet_tower::Tower2::<K>::zero();
eps.v = ab.v * be.v + ae.v * bb.v;
for i in 0..K {
eps.g[i] = ab.v * be.g[i] + ab.g[i] * be.v + ae.v * bb.g[i] + ae.g[i] * bb.v;
}
for i in 0..K {
for j in i..K {
let channel = ab.v * be.h[i][j]
+ ab.g[i] * be.g[j]
+ ab.g[j] * be.g[i]
+ ab.h[i][j] * be.v
+ ae.v * bb.h[i][j]
+ ae.g[i] * bb.g[j]
+ ae.g[j] * bb.g[i]
+ ae.h[i][j] * bb.v;
eps.h[i][j] = channel;
eps.h[j][i] = channel;
}
}
OneSeed {
base: self.base.mul(&o.base),
eps: Order2(eps),
}
}
fn neg(&self) -> Self {
OneSeed {
base: self.base.neg(),
eps: self.eps.neg(),
}
}
fn scale(&self, s: f64) -> Self {
OneSeed {
base: self.base.scale(s),
eps: self.eps.scale(s),
}
}
fn compose_unary(&self, d: [f64; 5]) -> Self {
let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
let b = &self.base.0;
let e = &self.eps.0;
let mut eps = crate::jet_tower::Tower2::<K>::zero();
eps.v = d[1] * e.v;
for i in 0..K {
eps.g[i] = d[2] * b.g[i] * e.v + d[1] * e.g[i];
}
for i in 0..K {
for j in i..K {
let channel = d[1] * e.h[i][j]
+ d[2] * (b.g[i] * e.g[j] + b.g[j] * e.g[i] + b.h[i][j] * e.v)
+ d[3] * b.g[i] * b.g[j] * e.v;
eps.h[i][j] = channel;
eps.h[j][i] = channel;
}
}
OneSeed {
base,
eps: Order2(eps),
}
}
fn constant_like(&self, v: f64) -> Self {
OneSeed {
base: self.base.constant_like(v),
eps: self.eps.constant_like(0.0),
}
}
fn with_value(&self, v: f64) -> Self {
OneSeed {
base: self.base.with_value(v),
eps: self.eps,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct OneSeedLane<L: Lane, const K: usize> {
pub base: Order2Lane<L, K>,
pub eps: Order2Lane<L, K>,
}
pub type OneSeedBatch<const K: usize> = OneSeedLane<wide::f64x4, K>;
impl<L: Lane, const K: usize> OneSeedLane<L, K> {
#[inline]
pub fn constant(c: L) -> Self {
OneSeedLane {
base: Order2Lane::constant(c),
eps: Order2Lane::constant(L::splat(0.0)),
}
}
#[inline]
pub fn variable(value: L, axis: usize) -> Self {
OneSeedLane {
base: Order2Lane::variable(value, axis),
eps: Order2Lane::constant(L::splat(0.0)),
}
}
#[inline]
pub fn seed_direction(value: L, axis: usize, u_axis: L) -> Self {
OneSeedLane {
base: Order2Lane::variable(value, axis),
eps: Order2Lane::constant(u_axis),
}
}
#[inline]
#[must_use]
pub fn contracted_third(&self) -> [[L; K]; K] {
self.eps.h
}
#[inline]
pub fn add(&self, o: &Self) -> Self {
OneSeedLane {
base: self.base.add(&o.base),
eps: self.eps.add(&o.eps),
}
}
#[inline]
pub fn sub(&self, o: &Self) -> Self {
OneSeedLane {
base: self.base.sub(&o.base),
eps: self.eps.sub(&o.eps),
}
}
#[inline]
pub fn mul(&self, o: &Self) -> Self {
let ab = &self.base;
let ae = &self.eps;
let bb = &o.base;
let be = &o.eps;
let mut eps = Order2Lane::constant(ab.v.mul(be.v).add(ae.v.mul(bb.v)));
for i in 0..K {
eps.g[i] =
ab.v.mul(be.g[i])
.add(ab.g[i].mul(be.v))
.add(ae.v.mul(bb.g[i]))
.add(ae.g[i].mul(bb.v));
}
for i in 0..K {
for j in i..K {
let channel =
ab.v.mul(be.h[i][j])
.add(ab.g[i].mul(be.g[j]))
.add(ab.g[j].mul(be.g[i]))
.add(ab.h[i][j].mul(be.v))
.add(ae.v.mul(bb.h[i][j]))
.add(ae.g[i].mul(bb.g[j]))
.add(ae.g[j].mul(bb.g[i]))
.add(ae.h[i][j].mul(bb.v));
eps.h[i][j] = channel;
eps.h[j][i] = channel;
}
}
OneSeedLane {
base: self.base.mul(&o.base),
eps,
}
}
#[inline]
pub fn neg(&self) -> Self {
OneSeedLane {
base: self.base.neg(),
eps: self.eps.neg(),
}
}
#[inline]
pub fn scale(&self, s: f64) -> Self {
OneSeedLane {
base: self.base.scale(s),
eps: self.eps.scale(s),
}
}
#[inline]
pub fn compose_unary(&self, d: [L; 5]) -> Self {
let base = self.base.compose_unary([d[0], d[1], d[2]]);
let b = &self.base;
let e = &self.eps;
let mut eps = Order2Lane::constant(d[1].mul(e.v));
for i in 0..K {
eps.g[i] = d[2].mul(b.g[i]).mul(e.v).add(d[1].mul(e.g[i]));
}
for i in 0..K {
for j in i..K {
let mixed = b.g[i]
.mul(e.g[j])
.add(b.g[j].mul(e.g[i]))
.add(b.h[i][j].mul(e.v));
let channel = d[1]
.mul(e.h[i][j])
.add(d[2].mul(mixed))
.add(d[3].mul(b.g[i]).mul(b.g[j]).mul(e.v));
eps.h[i][j] = channel;
eps.h[j][i] = channel;
}
}
OneSeedLane { base, eps }
}
#[inline]
pub fn exp(&self) -> Self {
let d = self.base.v.unary5(|u| {
let e = u.exp();
[e, e, e, e, e]
});
self.compose_unary(d)
}
#[inline]
pub fn ln(&self) -> Self {
let d = self.base.v.unary5(|u| {
let r = 1.0 / u;
[u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
});
self.compose_unary(d)
}
#[inline]
pub fn sqrt(&self) -> Self {
let d = self.base.v.unary5(|u| {
let s = u.sqrt();
[
s,
0.5 / s,
-0.25 / (u * s),
0.375 / (u * u * s),
-0.9375 / (u * u * u * s),
]
});
self.compose_unary(d)
}
#[inline]
pub fn recip(&self) -> Self {
let d = self.base.v.unary5(|u| {
let r = 1.0 / u;
let r2 = r * r;
[r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
});
self.compose_unary(d)
}
#[inline]
pub fn powf(&self, a: f64) -> Self {
let d = self.base.v.unary5(|u| {
[
u.powf(a),
a * u.powf(a - 1.0),
a * (a - 1.0) * u.powf(a - 2.0),
a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
]
});
self.compose_unary(d)
}
#[inline]
pub fn ln_gamma(&self) -> Self {
let d = self
.base
.v
.unary5(crate::jet_tower::ln_gamma_derivative_stack);
self.compose_unary(d)
}
#[inline]
pub fn digamma(&self) -> Self {
let d = self
.base
.v
.unary5(crate::jet_tower::digamma_derivative_stack);
self.compose_unary(d)
}
}
impl<const K: usize> OneSeedBatch<K> {
#[inline]
#[must_use]
pub fn lane(&self, i: usize) -> OneSeed<K> {
OneSeed {
base: self.base.lane(i),
eps: self.eps.lane(i),
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct TwoSeed<const K: usize> {
pub base: Order2<K>,
pub eps: Order2<K>,
pub del: Order2<K>,
pub eps_del: Order2<K>,
}
impl<const K: usize> TwoSeed<K> {
pub fn seed(x: f64, axis: usize, u_axis: f64, v_axis: f64) -> Self {
TwoSeed {
base: Order2::variable(x, axis),
eps: Order2::constant(u_axis),
del: Order2::constant(v_axis),
eps_del: Order2::constant(0.0),
}
}
pub fn contracted_fourth(&self) -> [[f64; K]; K] {
*self.eps_del.h()
}
}
impl<const K: usize> JetScalar<K> for TwoSeed<K> {
fn constant(c: f64) -> Self {
TwoSeed {
base: Order2::constant(c),
eps: Order2::constant(0.0),
del: Order2::constant(0.0),
eps_del: Order2::constant(0.0),
}
}
fn variable(x: f64, axis: usize) -> Self {
TwoSeed {
base: Order2::variable(x, axis),
eps: Order2::constant(0.0),
del: Order2::constant(0.0),
eps_del: Order2::constant(0.0),
}
}
}
impl<const K: usize> crate::nested_dual::JetField for TwoSeed<K> {
fn value(&self) -> f64 {
self.base.value()
}
fn add(&self, o: &Self) -> Self {
TwoSeed {
base: self.base.add(&o.base),
eps: self.eps.add(&o.eps),
del: self.del.add(&o.del),
eps_del: self.eps_del.add(&o.eps_del),
}
}
fn sub(&self, o: &Self) -> Self {
TwoSeed {
base: self.base.sub(&o.base),
eps: self.eps.sub(&o.eps),
del: self.del.sub(&o.del),
eps_del: self.eps_del.sub(&o.eps_del),
}
}
fn mul(&self, o: &Self) -> Self {
let a = self;
let b = o;
let base = a.base.mul(&b.base);
let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
let eps_del = a
.base
.mul(&b.eps_del)
.add(&a.eps.mul(&b.del))
.add(&a.del.mul(&b.eps))
.add(&a.eps_del.mul(&b.base));
TwoSeed {
base,
eps,
del,
eps_del,
}
}
fn neg(&self) -> Self {
TwoSeed {
base: self.base.neg(),
eps: self.eps.neg(),
del: self.del.neg(),
eps_del: self.eps_del.neg(),
}
}
fn scale(&self, s: f64) -> Self {
TwoSeed {
base: self.base.scale(s),
eps: self.eps.scale(s),
del: self.del.scale(s),
eps_del: self.eps_del.scale(s),
}
}
fn compose_unary(&self, d: [f64; 5]) -> Self {
let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]); let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]); let eps = fprime.mul(&self.eps);
let del = fprime.mul(&self.del);
let eps_del = fsecond
.mul(&self.eps)
.mul(&self.del)
.add(&fprime.mul(&self.eps_del));
TwoSeed {
base,
eps,
del,
eps_del,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct TwoSeedLane<L: Lane, const K: usize> {
pub base: Order2Lane<L, K>,
pub eps: Order2Lane<L, K>,
pub del: Order2Lane<L, K>,
pub eps_del: Order2Lane<L, K>,
}
pub type TwoSeedBatch<const K: usize> = TwoSeedLane<wide::f64x4, K>;
impl<L: Lane, const K: usize> TwoSeedLane<L, K> {
#[inline]
pub fn constant(c: L) -> Self {
let z = Order2Lane::constant(L::splat(0.0));
TwoSeedLane {
base: Order2Lane::constant(c),
eps: z,
del: z,
eps_del: z,
}
}
#[inline]
pub fn variable(value: L, axis: usize) -> Self {
let z = Order2Lane::constant(L::splat(0.0));
TwoSeedLane {
base: Order2Lane::variable(value, axis),
eps: z,
del: z,
eps_del: z,
}
}
#[inline]
pub fn seed(value: L, axis: usize, u_axis: L, v_axis: L) -> Self {
TwoSeedLane {
base: Order2Lane::variable(value, axis),
eps: Order2Lane::constant(u_axis),
del: Order2Lane::constant(v_axis),
eps_del: Order2Lane::constant(L::splat(0.0)),
}
}
#[inline]
#[must_use]
pub fn contracted_fourth(&self) -> [[L; K]; K] {
self.eps_del.h
}
#[inline]
pub fn add(&self, o: &Self) -> Self {
TwoSeedLane {
base: self.base.add(&o.base),
eps: self.eps.add(&o.eps),
del: self.del.add(&o.del),
eps_del: self.eps_del.add(&o.eps_del),
}
}
#[inline]
pub fn sub(&self, o: &Self) -> Self {
TwoSeedLane {
base: self.base.sub(&o.base),
eps: self.eps.sub(&o.eps),
del: self.del.sub(&o.del),
eps_del: self.eps_del.sub(&o.eps_del),
}
}
#[inline]
pub fn mul(&self, o: &Self) -> Self {
let a = self;
let b = o;
let base = a.base.mul(&b.base);
let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
let eps_del = a
.base
.mul(&b.eps_del)
.add(&a.eps.mul(&b.del))
.add(&a.del.mul(&b.eps))
.add(&a.eps_del.mul(&b.base));
TwoSeedLane {
base,
eps,
del,
eps_del,
}
}
#[inline]
pub fn neg(&self) -> Self {
TwoSeedLane {
base: self.base.neg(),
eps: self.eps.neg(),
del: self.del.neg(),
eps_del: self.eps_del.neg(),
}
}
#[inline]
pub fn scale(&self, s: f64) -> Self {
TwoSeedLane {
base: self.base.scale(s),
eps: self.eps.scale(s),
del: self.del.scale(s),
eps_del: self.eps_del.scale(s),
}
}
#[inline]
pub fn compose_unary(&self, d: [L; 5]) -> Self {
let base = self.base.compose_unary([d[0], d[1], d[2]]);
let fprime = self.base.compose_unary([d[1], d[2], d[3]]);
let fsecond = self.base.compose_unary([d[2], d[3], d[4]]);
let eps = fprime.mul(&self.eps);
let del = fprime.mul(&self.del);
let eps_del = fsecond
.mul(&self.eps)
.mul(&self.del)
.add(&fprime.mul(&self.eps_del));
TwoSeedLane {
base,
eps,
del,
eps_del,
}
}
#[inline]
pub fn exp(&self) -> Self {
let d = self.base.v.unary5(|u| {
let e = u.exp();
[e, e, e, e, e]
});
self.compose_unary(d)
}
#[inline]
pub fn ln(&self) -> Self {
let d = self.base.v.unary5(|u| {
let r = 1.0 / u;
[u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
});
self.compose_unary(d)
}
#[inline]
pub fn sqrt(&self) -> Self {
let d = self.base.v.unary5(|u| {
let s = u.sqrt();
[
s,
0.5 / s,
-0.25 / (u * s),
0.375 / (u * u * s),
-0.9375 / (u * u * u * s),
]
});
self.compose_unary(d)
}
#[inline]
pub fn recip(&self) -> Self {
let d = self.base.v.unary5(|u| {
let r = 1.0 / u;
let r2 = r * r;
[r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
});
self.compose_unary(d)
}
#[inline]
pub fn powf(&self, a: f64) -> Self {
let d = self.base.v.unary5(|u| {
[
u.powf(a),
a * u.powf(a - 1.0),
a * (a - 1.0) * u.powf(a - 2.0),
a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
]
});
self.compose_unary(d)
}
#[inline]
pub fn ln_gamma(&self) -> Self {
let d = self
.base
.v
.unary5(crate::jet_tower::ln_gamma_derivative_stack);
self.compose_unary(d)
}
#[inline]
pub fn digamma(&self) -> Self {
let d = self
.base
.v
.unary5(crate::jet_tower::digamma_derivative_stack);
self.compose_unary(d)
}
}
impl<const K: usize> TwoSeedBatch<K> {
#[inline]
#[must_use]
pub fn lane(&self, i: usize) -> TwoSeed<K> {
TwoSeed {
base: self.base.lane(i),
eps: self.eps.lane(i),
del: self.del.lane(i),
eps_del: self.eps_del.lane(i),
}
}
}
impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower3<K> {
fn constant(c: f64) -> Self {
crate::jet_tower::Tower3::constant(c)
}
fn variable(x: f64, axis: usize) -> Self {
crate::jet_tower::Tower3::variable(x, axis)
}
}
impl<const K: usize> crate::nested_dual::JetField for crate::jet_tower::Tower3<K> {
fn value(&self) -> f64 {
self.v
}
fn add(&self, o: &Self) -> Self {
*self + *o
}
fn sub(&self, o: &Self) -> Self {
*self + o.scale(-1.0)
}
fn mul(&self, o: &Self) -> Self {
crate::jet_tower::Tower3::mul(self, o)
}
fn neg(&self) -> Self {
self.scale(-1.0)
}
fn scale(&self, s: f64) -> Self {
crate::jet_tower::Tower3::scale(self, s)
}
fn compose_unary(&self, d: [f64; 5]) -> Self {
crate::jet_tower::Tower3::compose_unary(self, [d[0], d[1], d[2], d[3]])
}
}
impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower4<K> {
fn constant(c: f64) -> Self {
crate::jet_tower::Tower4::constant(c)
}
fn variable(x: f64, axis: usize) -> Self {
crate::jet_tower::Tower4::variable(x, axis)
}
}
impl<const K: usize> crate::nested_dual::JetField for crate::jet_tower::Tower4<K> {
fn value(&self) -> f64 {
self.v
}
fn add(&self, o: &Self) -> Self {
*self + *o
}
fn sub(&self, o: &Self) -> Self {
*self - *o
}
fn mul(&self, o: &Self) -> Self {
crate::jet_tower::Tower4::mul(self, o)
}
fn neg(&self) -> Self {
self.scale(-1.0)
}
fn scale(&self, s: f64) -> Self {
crate::jet_tower::Tower4::scale(self, s)
}
fn compose_unary(&self, d: [f64; 5]) -> Self {
crate::jet_tower::Tower4::compose_unary(self, d)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::jet_tower::{RowProgram, Tower4, program_full_tower};
use crate::nested_dual::JetField;
struct DenseSymmetric3([[f64; 3]; 3]);
impl SymmetricQuadraticCoefficients for DenseSymmetric3 {
fn dimension(&self) -> usize {
3
}
fn multiply(&self, input: &[f64], output: &mut [f64]) {
assert_eq!(input.len(), 3);
assert_eq!(output.len(), 3);
for (row, output) in output.iter_mut().enumerate() {
*output = (0..3)
.map(|column| self.0[row][column] * input[column])
.sum();
}
}
fn coefficient(&self, row: usize, column: usize) -> f64 {
self.0[row][column]
}
}
#[test]
fn symmetric_quadratic_order2_lowerings_match_scalar_program() {
const K: usize = 4;
let coefficients = DenseSymmetric3([[1.2, 0.3, -0.2], [0.3, 0.8, 0.15], [-0.2, 0.15, 1.5]]);
let values = [0.4, -0.7, 1.1, 0.25];
let fixed_vars: [Order2<K>; K] =
std::array::from_fn(|axis| Order2::variable(values[axis], axis));
let fixed_inputs = [
fixed_vars[0].mul(&fixed_vars[1]).add(&fixed_vars[3]),
fixed_vars[1].exp().add(&fixed_vars[2].scale(0.4)),
fixed_vars[2].mul(&fixed_vars[2]).sub(&fixed_vars[0]),
];
let fixed_direct = Order2::symmetric_quadratic_form(&fixed_inputs, &coefficients);
let fixed_scalar = symmetric_quadratic_form_default(
&fixed_inputs,
&coefficients,
Order2::constant,
JetField::add,
JetField::mul,
JetField::scale,
);
let weights = [0.7, -1.1, 0.35];
let fixed_linear_direct = Order2::linear_combination(&fixed_inputs, &weights);
let fixed_linear_scalar = linear_combination_default(
&fixed_inputs,
&weights,
Order2::constant,
JetField::add,
JetField::scale,
);
let derivative_stacks = [
[0.8, -0.3, 0.7, 0.0, 0.0],
[-0.2, 1.1, -0.4, 0.0, 0.0],
[1.4, 0.25, 0.6, 0.0, 0.0],
];
let fixed_add_direct = fixed_inputs[0].add_constant(0.65);
let fixed_add_scalar = fixed_inputs[0].add(&Order2::constant(0.65));
let fixed_multiply_add_direct =
fixed_inputs[0].multiply_add(&fixed_inputs[1], &fixed_inputs[2]);
let fixed_multiply_add_scalar = multiply_add_default(
&fixed_inputs[0],
&fixed_inputs[1],
&fixed_inputs[2],
JetField::mul,
JetField::add,
);
let fixed_composed_direct = Order2::composed_sum(&fixed_inputs, &derivative_stacks);
let fixed_composed_scalar = composed_sum_default(
&fixed_inputs,
&derivative_stacks,
Order2::constant,
JetField::add,
JetField::compose_unary,
);
let value_vars: [RuntimeValue; K] =
std::array::from_fn(|axis| RuntimeValue::variable(values[axis], axis, K, &()));
let value_inputs = [
value_vars[0].mul(&value_vars[1]).add(&value_vars[3]),
value_vars[1].exp().add(&value_vars[2].scale(0.4)),
value_vars[2].mul(&value_vars[2]).sub(&value_vars[0]),
];
let value_quadratic =
RuntimeValue::symmetric_quadratic_form(&value_inputs, &coefficients, K, &());
let value_linear = RuntimeValue::linear_combination(&value_inputs, &weights, K, &());
let value_composed = RuntimeValue::composed_sum(&value_inputs, &derivative_stacks, K, &());
let arena = DynamicJetArena::new();
let dynamic_vars: [DynamicOrder2<'_>; K] =
std::array::from_fn(|axis| DynamicOrder2::variable(values[axis], axis, K, &arena));
let dynamic_inputs = [
dynamic_vars[0].mul(&dynamic_vars[1]).add(&dynamic_vars[3]),
dynamic_vars[1].exp().add(&dynamic_vars[2].scale(0.4)),
dynamic_vars[2].mul(&dynamic_vars[2]).sub(&dynamic_vars[0]),
];
let dynamic_direct =
DynamicOrder2::symmetric_quadratic_form(&dynamic_inputs, &coefficients, K, &arena);
let dynamic_scalar = symmetric_quadratic_form_default(
&dynamic_inputs,
&coefficients,
|value| DynamicOrder2::constant(value, K, &arena),
RuntimeJetScalar::add,
RuntimeJetScalar::mul,
RuntimeJetScalar::scale,
);
let dynamic_linear_direct =
DynamicOrder2::linear_combination(&dynamic_inputs, &weights, K, &arena);
let dynamic_linear_scalar = linear_combination_default(
&dynamic_inputs,
&weights,
|value| DynamicOrder2::constant(value, K, &arena),
RuntimeJetScalar::add,
RuntimeJetScalar::scale,
);
let dynamic_add_direct = dynamic_inputs[0].add_constant(0.65);
let dynamic_add_scalar = dynamic_inputs[0].add(&DynamicOrder2::constant(0.65, K, &arena));
let dynamic_multiply_add_direct =
dynamic_inputs[0].multiply_add(&dynamic_inputs[1], &dynamic_inputs[2]);
let dynamic_multiply_add_scalar = multiply_add_default(
&dynamic_inputs[0],
&dynamic_inputs[1],
&dynamic_inputs[2],
RuntimeJetScalar::mul,
RuntimeJetScalar::add,
);
let dynamic_composed_direct =
DynamicOrder2::composed_sum(&dynamic_inputs, &derivative_stacks, K, &arena);
let dynamic_composed_scalar = composed_sum_default(
&dynamic_inputs,
&derivative_stacks,
|value| DynamicOrder2::constant(value, K, &arena),
RuntimeJetScalar::add,
RuntimeJetScalar::compose_unary,
);
let tolerance = 2.0e-13;
for (label, actual, expected) in [
("fixed value", fixed_direct.value(), fixed_scalar.value()),
(
"zero-order quadratic value",
value_quadratic.value(),
fixed_direct.value(),
),
(
"zero-order linear value",
value_linear.value(),
fixed_linear_direct.value(),
),
(
"zero-order composed value",
value_composed.value(),
fixed_composed_direct.value(),
),
(
"dynamic value",
dynamic_direct.value(),
dynamic_scalar.value(),
),
(
"fixed linear value",
fixed_linear_direct.value(),
fixed_linear_scalar.value(),
),
(
"dynamic linear value",
dynamic_linear_direct.value(),
dynamic_linear_scalar.value(),
),
(
"fixed add-constant value",
fixed_add_direct.value(),
fixed_add_scalar.value(),
),
(
"dynamic add-constant value",
dynamic_add_direct.value(),
dynamic_add_scalar.value(),
),
(
"fixed multiply-add value",
fixed_multiply_add_direct.value(),
fixed_multiply_add_scalar.value(),
),
(
"dynamic multiply-add value",
dynamic_multiply_add_direct.value(),
dynamic_multiply_add_scalar.value(),
),
(
"fixed composed-sum value",
fixed_composed_direct.value(),
fixed_composed_scalar.value(),
),
(
"dynamic composed-sum value",
dynamic_composed_direct.value(),
dynamic_composed_scalar.value(),
),
] {
assert!(
(actual - expected).abs() <= tolerance * actual.abs().max(expected.abs()).max(1.0),
"{label}: direct={actual:+.16e}, scalar={expected:+.16e}"
);
}
for primary_a in 0..K {
for (label, actual, expected) in [
(
"fixed gradient",
fixed_direct.g()[primary_a],
fixed_scalar.g()[primary_a],
),
(
"dynamic gradient",
dynamic_direct.g()[primary_a],
dynamic_scalar.g()[primary_a],
),
(
"fixed linear gradient",
fixed_linear_direct.g()[primary_a],
fixed_linear_scalar.g()[primary_a],
),
(
"dynamic linear gradient",
dynamic_linear_direct.g()[primary_a],
dynamic_linear_scalar.g()[primary_a],
),
(
"fixed add-constant gradient",
fixed_add_direct.g()[primary_a],
fixed_add_scalar.g()[primary_a],
),
(
"dynamic add-constant gradient",
dynamic_add_direct.g()[primary_a],
dynamic_add_scalar.g()[primary_a],
),
(
"fixed multiply-add gradient",
fixed_multiply_add_direct.g()[primary_a],
fixed_multiply_add_scalar.g()[primary_a],
),
(
"dynamic multiply-add gradient",
dynamic_multiply_add_direct.g()[primary_a],
dynamic_multiply_add_scalar.g()[primary_a],
),
(
"fixed composed-sum gradient",
fixed_composed_direct.g()[primary_a],
fixed_composed_scalar.g()[primary_a],
),
(
"dynamic composed-sum gradient",
dynamic_composed_direct.g()[primary_a],
dynamic_composed_scalar.g()[primary_a],
),
] {
assert!(
(actual - expected).abs()
<= tolerance * actual.abs().max(expected.abs()).max(1.0),
"{label}[{primary_a}]: direct={actual:+.16e}, scalar={expected:+.16e}"
);
}
for primary_b in 0..K {
for (label, actual, expected) in [
(
"fixed Hessian",
fixed_direct.h()[primary_a][primary_b],
fixed_scalar.h()[primary_a][primary_b],
),
(
"dynamic Hessian",
dynamic_direct.h_at(primary_a, primary_b),
dynamic_scalar.h_at(primary_a, primary_b),
),
(
"fixed linear Hessian",
fixed_linear_direct.h()[primary_a][primary_b],
fixed_linear_scalar.h()[primary_a][primary_b],
),
(
"dynamic linear Hessian",
dynamic_linear_direct.h_at(primary_a, primary_b),
dynamic_linear_scalar.h_at(primary_a, primary_b),
),
(
"fixed add-constant Hessian",
fixed_add_direct.h()[primary_a][primary_b],
fixed_add_scalar.h()[primary_a][primary_b],
),
(
"dynamic add-constant Hessian",
dynamic_add_direct.h_at(primary_a, primary_b),
dynamic_add_scalar.h_at(primary_a, primary_b),
),
(
"fixed multiply-add Hessian",
fixed_multiply_add_direct.h()[primary_a][primary_b],
fixed_multiply_add_scalar.h()[primary_a][primary_b],
),
(
"dynamic multiply-add Hessian",
dynamic_multiply_add_direct.h_at(primary_a, primary_b),
dynamic_multiply_add_scalar.h_at(primary_a, primary_b),
),
(
"fixed composed-sum Hessian",
fixed_composed_direct.h()[primary_a][primary_b],
fixed_composed_scalar.h()[primary_a][primary_b],
),
(
"dynamic composed-sum Hessian",
dynamic_composed_direct.h_at(primary_a, primary_b),
dynamic_composed_scalar.h_at(primary_a, primary_b),
),
] {
assert!(
(actual - expected).abs()
<= tolerance * actual.abs().max(expected.abs()).max(1.0),
"{label}[{primary_a},{primary_b}]: direct={actual:+.16e}, scalar={expected:+.16e}"
);
}
}
}
}
#[test]
fn compiled_product_affine_and_fused_nodes_match_scalar_program_randomized() {
const K: usize = 4;
const TERMS: usize = 10;
fn sample(state: &mut u64) -> f64 {
*state ^= *state << 13;
*state ^= *state >> 7;
*state ^= *state << 17;
let unit = (*state >> 11) as f64 * (1.0 / ((1_u64 << 53) as f64));
2.0 * unit - 1.0
}
fn arbitrary_order2<const K: usize>(state: &mut u64) -> Order2<K> {
let mut tower = crate::jet_tower::Tower2::zero();
tower.v = sample(state);
for primary in 0..K {
tower.g[primary] = sample(state);
for other in primary..K {
let channel = sample(state);
tower.h[primary][other] = channel;
tower.h[other][primary] = channel;
}
}
Order2(tower)
}
fn close(actual: f64, expected: f64, case: usize, label: &str) {
let tolerance = 2.0e-12 * actual.abs().max(expected.abs()).max(1.0);
assert!(
(actual - expected).abs() <= tolerance,
"case {case} {label}: direct={actual:+.16e}, scalar={expected:+.16e}, tolerance={tolerance:.3e}"
);
}
let mut state = 0x932a_ff1e_c0de_5eed_u64;
for case in 0..256 {
let fixed_inputs: [Order2<K>; TERMS] =
std::array::from_fn(|_| arbitrary_order2(&mut state));
let mut input_scales: [f64; TERMS] = std::array::from_fn(|_| sample(&mut state));
input_scales[0] = -1.25;
input_scales[1] = 0.0;
input_scales[4] = 1.25;
let addend_scales: [f64; TERMS] = std::array::from_fn(|term| match term % 4 {
0 => 0.0,
1 => 1.0,
2 => -0.75,
_ => 0.35,
});
let derivative_stacks: [[f64; 5]; TERMS] =
std::array::from_fn(|_| std::array::from_fn(|_| sample(&mut state)));
let input_shift = sample(&mut state);
let mut fixed_lefts = std::array::from_fn(|term| &fixed_inputs[term]);
fixed_lefts[4] = &fixed_inputs[0];
fixed_lefts[5] = &fixed_inputs[0];
let fixed_right = &fixed_inputs[1];
let fixed_addend = &fixed_inputs[2];
let fixed_product_direct = fixed_inputs[0].product(&fixed_inputs[1]);
let fixed_product_scalar = fixed_inputs[0].mul(&fixed_inputs[1]);
let fixed_affine_direct =
fixed_inputs[2].affine_compose(input_scales[2], input_shift, derivative_stacks[2]);
let fixed_affine_scalar = affine_compose_default(
&fixed_inputs[2],
input_scales[2],
input_shift,
derivative_stacks[2],
JetField::scale,
Order2::add_constant,
JetField::compose_unary,
);
let fixed_sum_direct =
Order2::affine_composed_sum(&fixed_inputs, &input_scales, &derivative_stacks);
let fixed_sum_scalar = affine_composed_sum_default(
&fixed_inputs,
&input_scales,
&derivative_stacks,
Order2::constant,
JetField::add,
JetField::scale,
Order2::add_constant,
JetField::compose_unary,
);
let fixed_fused_direct = Order2::shared_multiply_add_affine_composed_sum(
&fixed_lefts,
fixed_right,
fixed_addend,
&addend_scales,
&input_scales,
&derivative_stacks,
);
let fixed_fused_scalar = shared_multiply_add_affine_composed_sum_default(
&fixed_lefts,
fixed_right,
fixed_addend,
&addend_scales,
&input_scales,
&derivative_stacks,
Order2::constant,
JetField::add,
JetField::mul,
JetField::scale,
Order2::multiply_add,
Order2::affine_compose,
);
let arena = DynamicJetArena::new();
let dynamic_inputs: [DynamicOrder2<'_>; TERMS] = std::array::from_fn(|term| {
DynamicOrder2::from_channel_functions(
fixed_inputs[term].value(),
K,
&arena,
|primary| fixed_inputs[term].g()[primary],
|primary, other| fixed_inputs[term].h()[primary][other],
)
});
let dynamic_product_direct = dynamic_inputs[0].product(&dynamic_inputs[1]);
let dynamic_product_scalar = dynamic_inputs[0].mul(&dynamic_inputs[1]);
let dynamic_affine_direct = dynamic_inputs[2].affine_compose(
input_scales[2],
input_shift,
derivative_stacks[2],
);
let dynamic_affine_scalar = affine_compose_default(
&dynamic_inputs[2],
input_scales[2],
input_shift,
derivative_stacks[2],
RuntimeJetScalar::scale,
|input, constant| input.add_constant(constant),
RuntimeJetScalar::compose_unary,
);
let dynamic_sum_direct = DynamicOrder2::affine_composed_sum(
&dynamic_inputs,
&input_scales,
&derivative_stacks,
K,
&arena,
);
let dynamic_sum_scalar = affine_composed_sum_default(
&dynamic_inputs,
&input_scales,
&derivative_stacks,
|value| DynamicOrder2::constant(value, K, &arena),
RuntimeJetScalar::add,
RuntimeJetScalar::scale,
|input, constant| input.add_constant(constant),
RuntimeJetScalar::compose_unary,
);
let mut dynamic_lefts = std::array::from_fn(|term| &dynamic_inputs[term]);
dynamic_lefts[4] = &dynamic_inputs[0];
dynamic_lefts[5] = &dynamic_inputs[0];
let dynamic_right = &dynamic_inputs[1];
let dynamic_addend = &dynamic_inputs[2];
let dynamic_fused_direct = DynamicOrder2::shared_multiply_add_affine_composed_sum(
&dynamic_lefts,
dynamic_right,
dynamic_addend,
&addend_scales,
&input_scales,
&derivative_stacks,
K,
&arena,
);
let dynamic_fused_scalar = shared_multiply_add_affine_composed_sum_default(
&dynamic_lefts,
dynamic_right,
dynamic_addend,
&addend_scales,
&input_scales,
&derivative_stacks,
|value| DynamicOrder2::constant(value, K, &arena),
RuntimeJetScalar::add,
RuntimeJetScalar::mul,
RuntimeJetScalar::scale,
RuntimeJetScalar::multiply_add,
|input, scale, shift, stack| input.affine_compose(scale, shift, stack),
);
for (label, actual, expected) in [
(
"fixed product value",
fixed_product_direct.value(),
fixed_product_scalar.value(),
),
(
"fixed affine value",
fixed_affine_direct.value(),
fixed_affine_scalar.value(),
),
(
"fixed affine sum value",
fixed_sum_direct.value(),
fixed_sum_scalar.value(),
),
(
"fixed fused value",
fixed_fused_direct.value(),
fixed_fused_scalar.value(),
),
(
"dynamic product value",
dynamic_product_direct.value(),
dynamic_product_scalar.value(),
),
(
"dynamic affine value",
dynamic_affine_direct.value(),
dynamic_affine_scalar.value(),
),
(
"dynamic affine sum value",
dynamic_sum_direct.value(),
dynamic_sum_scalar.value(),
),
(
"dynamic fused value",
dynamic_fused_direct.value(),
dynamic_fused_scalar.value(),
),
] {
close(actual, expected, case, label);
}
for primary in 0..K {
for (label, actual, expected) in [
(
"fixed product gradient",
fixed_product_direct.g()[primary],
fixed_product_scalar.g()[primary],
),
(
"fixed affine gradient",
fixed_affine_direct.g()[primary],
fixed_affine_scalar.g()[primary],
),
(
"fixed affine sum gradient",
fixed_sum_direct.g()[primary],
fixed_sum_scalar.g()[primary],
),
(
"fixed fused gradient",
fixed_fused_direct.g()[primary],
fixed_fused_scalar.g()[primary],
),
(
"dynamic product gradient",
dynamic_product_direct.g()[primary],
dynamic_product_scalar.g()[primary],
),
(
"dynamic affine gradient",
dynamic_affine_direct.g()[primary],
dynamic_affine_scalar.g()[primary],
),
(
"dynamic affine sum gradient",
dynamic_sum_direct.g()[primary],
dynamic_sum_scalar.g()[primary],
),
(
"dynamic fused gradient",
dynamic_fused_direct.g()[primary],
dynamic_fused_scalar.g()[primary],
),
] {
close(actual, expected, case, label);
}
for other in 0..K {
for (label, actual, expected) in [
(
"fixed product Hessian",
fixed_product_direct.h()[primary][other],
fixed_product_scalar.h()[primary][other],
),
(
"fixed affine Hessian",
fixed_affine_direct.h()[primary][other],
fixed_affine_scalar.h()[primary][other],
),
(
"fixed affine sum Hessian",
fixed_sum_direct.h()[primary][other],
fixed_sum_scalar.h()[primary][other],
),
(
"fixed fused Hessian",
fixed_fused_direct.h()[primary][other],
fixed_fused_scalar.h()[primary][other],
),
(
"dynamic product Hessian",
dynamic_product_direct.h_at(primary, other),
dynamic_product_scalar.h_at(primary, other),
),
(
"dynamic affine Hessian",
dynamic_affine_direct.h_at(primary, other),
dynamic_affine_scalar.h_at(primary, other),
),
(
"dynamic affine sum Hessian",
dynamic_sum_direct.h_at(primary, other),
dynamic_sum_scalar.h_at(primary, other),
),
(
"dynamic fused Hessian",
dynamic_fused_direct.h_at(primary, other),
dynamic_fused_scalar.h_at(primary, other),
),
] {
close(actual, expected, case, label);
}
}
}
}
}
#[test]
fn shared_product_composition_accepts_empty_expression() {
const K: usize = 4;
let fixed_terms: [&Order2<K>; 0] = [];
let fixed_shared = Order2::constant(1.0);
let scales: [f64; 0] = [];
let stacks: [[f64; 5]; 0] = [];
let fixed = Order2::shared_multiply_add_affine_composed_sum(
&fixed_terms,
&fixed_shared,
&fixed_shared,
&scales,
&scales,
&stacks,
);
assert_eq!(fixed.value().to_bits(), 0.0_f64.to_bits());
assert!(fixed.g().iter().all(|&channel| channel == 0.0));
assert!(fixed.h().iter().flatten().all(|&channel| channel == 0.0));
let value_terms: [&RuntimeValue; 0] = [];
let value_shared = RuntimeValue::constant(1.0, K, &());
let value = RuntimeValue::shared_multiply_add_affine_composed_sum(
&value_terms,
&value_shared,
&value_shared,
&scales,
&scales,
&stacks,
K,
&(),
);
assert_eq!(value.value().to_bits(), 0.0_f64.to_bits());
assert_eq!(value.dimension(), K);
let arena = DynamicJetArena::new();
let dynamic_terms: [&DynamicOrder2<'_>; 0] = [];
let dynamic_shared = DynamicOrder2::constant(1.0, K, &arena);
let dynamic = DynamicOrder2::shared_multiply_add_affine_composed_sum(
&dynamic_terms,
&dynamic_shared,
&dynamic_shared,
&scales,
&scales,
&stacks,
K,
&arena,
);
assert_eq!(dynamic.value().to_bits(), 0.0_f64.to_bits());
assert!((0..K).all(|axis| dynamic.g()[axis] == 0.0));
assert!((0..K).all(|row| (0..K).all(|column| dynamic.h_at(row, column) == 0.0)));
}
#[test]
fn runtime_fused_product_composition_preserves_tower4_channels() {
const K: usize = 2;
const N: usize = 9;
let vars = [
Tower4::<K>::variable(0.37, 0),
Tower4::<K>::variable(-0.61, 1),
];
let upstream = [
vars[0].mul(&vars[1]).add(&vars[0].exp()),
vars[1].mul(&vars[1]).add(&vars[0].scale(0.3)),
vars[0].mul(&vars[0]).sub(&vars[1].scale(-0.2)),
];
let mut lefts: [Tower4<K>; N] = std::array::from_fn(|term| upstream[term % upstream.len()]);
lefts[4] = lefts[0];
lefts[5] = lefts[0];
let right = upstream[1];
let addend = upstream[2];
let addend_scales: [f64; N] = std::array::from_fn(|term| [-0.0, 1.0, -0.7, 0.25][term % 4]);
let mut input_scales: [f64; N] =
std::array::from_fn(|term| [0.0, -1.3, 0.45, 1.1][term % 4]);
input_scales[0] = -1.1;
input_scales[4] = 1.1;
let stacks: [[f64; 5]; N] = std::array::from_fn(|term| {
let t = term as f64 + 1.0;
[0.17 * t, -0.11 * t, 0.07 * t, -0.03 * t, 0.013 * t]
});
let expected = (0..N).fold(Tower4::<K>::constant(0.0), |sum, term| {
let inner = if addend_scales[term] == 0.0 {
lefts[term].mul(&right)
} else if addend_scales[term] == 1.0 {
JetScalar::multiply_add(&lefts[term], &right, &addend)
} else {
JetScalar::multiply_add(&lefts[term], &right, &addend.scale(addend_scales[term]))
};
sum.add(&JetScalar::affine_compose(
&inner,
input_scales[term],
0.0,
stacks[term],
))
});
let wrapped_lefts: [FixedRuntimeJet<Tower4<K>, K>; N] =
std::array::from_fn(|term| FixedRuntimeJet::from_inner(lefts[term]));
let wrapped_right = FixedRuntimeJet::from_inner(right);
let wrapped_addend = FixedRuntimeJet::from_inner(addend);
let mut wrapped_left_refs: [&FixedRuntimeJet<Tower4<K>, K>; N] =
std::array::from_fn(|term| &wrapped_lefts[term]);
wrapped_left_refs[4] = &wrapped_lefts[0];
wrapped_left_refs[5] = &wrapped_lefts[0];
let actual = FixedRuntimeJet::<Tower4<K>, K>::shared_multiply_add_affine_composed_sum(
&wrapped_left_refs,
&wrapped_right,
&wrapped_addend,
&addend_scales,
&input_scales,
&stacks,
K,
&(),
)
.into_inner();
let same = |label: &str, got: f64, want: f64| {
let tolerance = 2.0e-13 * got.abs().max(want.abs()).max(1.0);
assert!(
(got - want).abs() <= tolerance,
"{label}: got={got:+.17e}, want={want:+.17e}, tolerance={tolerance:.3e}"
);
};
same("value", actual.v, expected.v);
for a in 0..K {
same("gradient", actual.g[a], expected.g[a]);
for b in 0..K {
same("Hessian", actual.h[a][b], expected.h[a][b]);
for c in 0..K {
same("third", actual.t3[a][b][c], expected.t3[a][b][c]);
for d in 0..K {
same("fourth", actual.t4[a][b][c][d], expected.t4[a][b][c][d]);
}
}
}
}
}
fn row_expr<S: JetScalar<2>>(p: &[S; 2]) -> S {
let g = p[0].mul(&p[1]).exp();
let inner = g.add(&S::constant(2.0));
let radic = p[0].mul(&p[0]).add(&S::constant(1.0)).sqrt();
inner.mul(&radic).sub(&p[1].mul(&p[1]).scale(0.5))
}
struct ExprProgram {
p: [f64; 2],
}
impl RowProgram<2> for ExprProgram {
fn n_rows(&self) -> usize {
1
}
fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
if row >= self.n_rows() {
return Err(format!("ExprProgram: row {row} out of range"));
}
Ok(self.p)
}
fn eval<S: JetScalar<2>>(&self, row: usize, p: &[S; 2]) -> Result<S, String> {
if row >= self.n_rows() {
return Err(format!("ExprProgram: row {row} out of range"));
}
Ok(row_expr(p))
}
}
const SEED: [f64; 2] = [0.37, -0.81];
const U: [f64; 2] = [0.6, -0.2];
const V: [f64; 2] = [-0.4, 1.1];
const TOL: f64 = 1e-10;
fn close(a: f64, b: f64, label: &str) {
let band = TOL + TOL * a.abs().max(b.abs());
assert!(
(a - b).abs() <= band,
"{label}: {a:+.15e} vs {b:+.15e} (band {band:.3e})"
);
}
fn tower() -> Tower4<2> {
*program_full_tower(&ExprProgram { p: SEED }, 0).expect("tower")
}
#[test]
fn order2_matches_tower_value_grad_hessian() {
let t = tower();
let vars: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
let s = row_expr(&vars);
close(s.value(), t.v, "value");
for a in 0..2 {
close(s.0.g[a], t.g[a], &format!("grad[{a}]"));
for b in 0..2 {
close(s.h()[a][b], t.h[a][b], &format!("hess[{a}][{b}]"));
}
}
}
#[test]
fn mapped_order2_accumulator_matches_dense_overlapping_atoms() {
const K: usize = 4;
let p = [0.2_f64, 0.7, -0.4, 0.3];
let dense_vars: [Order2<K>; K] =
std::array::from_fn(|axis| Order2::variable(p[axis], axis));
let dense_q0 = dense_vars[3].mul(&dense_vars[1]).add(&dense_vars[3].exp());
let dense_q1 = dense_vars[1].mul(&dense_vars[2]).sub(&dense_vars[2]);
let dense = dense_q0.ln().add(&dense_q1.exp());
let local_q0_vars: [Order2<2>; 2] =
std::array::from_fn(|axis| Order2::variable(p[[3, 1][axis]], axis));
let local_q0 = local_q0_vars[0]
.mul(&local_q0_vars[1])
.add(&local_q0_vars[0].exp());
let local_q1_vars: [Order2<2>; 2] =
std::array::from_fn(|axis| Order2::variable(p[[1, 2][axis]], axis));
let local_q1 = local_q1_vars[0]
.mul(&local_q1_vars[1])
.sub(&local_q1_vars[1]);
let q0 = local_q0.value();
let q1_exp = local_q1.value().exp();
let mut lowered = MappedOrder2Accumulator::<K>::zero();
lowered.add_composed(
&local_q0,
[3, 1],
[q0.ln(), q0.recip(), -1.0 / (q0 * q0)],
false,
[false, false],
[false, false, false],
);
lowered.add_composed(
&local_q1,
[1, 2],
[q1_exp, q1_exp, q1_exp],
true,
[true, false],
[true, false, false],
);
let (value, gradient, hessian) = lowered.into_channels();
close(value, dense.value(), "mapped value");
for i in 0..K {
close(gradient[i], dense.g()[i], &format!("mapped gradient[{i}]"));
for j in 0..K {
close(
hessian[i][j],
dense.h()[i][j],
&format!("mapped Hessian[{i},{j}]"),
);
}
}
}
#[test]
#[should_panic(expected = "mapped atom axes must be injective")]
fn mapped_order2_accumulator_rejects_duplicate_axes() {
let vars: [Order2<2>; 2] = std::array::from_fn(|axis| Order2::variable(0.2, axis));
let atom = vars[0].add(&vars[1]);
let mut lowered = MappedOrder2Accumulator::<2>::zero();
lowered.add_composed(
&atom,
[1, 1],
[0.4, 1.0, 0.0],
false,
[false, false],
[false, false, false],
);
}
#[test]
#[should_panic(expected = "mapped atom axis must be within")]
fn mapped_order2_accumulator_rejects_out_of_range_axes() {
let atom = Order2::<1>::variable(0.2, 0);
let mut lowered = MappedOrder2Accumulator::<2>::zero();
lowered.add_composed(&atom, [2], [0.2, 1.0, 0.0], false, [false], [false]);
}
#[test]
fn dynamic_order2_accumulator_matches_dense_composed_sum() {
const K: usize = 4;
struct Term {
first: f64,
second: f64,
gradient: [f64; K],
hessian: [[f64; K]; K],
}
impl DynamicOrder2Term for Term {
fn outer_first(&self) -> f64 {
self.first
}
fn outer_second(&self) -> f64 {
self.second
}
fn inner_gradient(&self, axis: usize) -> f64 {
self.gradient[axis]
}
fn inner_hessian(&self, row: usize, column: usize) -> f64 {
self.hessian[row][column]
}
}
let p = [0.7, -0.3, 0.2, 0.8];
let vars: [Order2<K>; K] = std::array::from_fn(|axis| Order2::variable(p[axis], axis));
let first_atom = vars[0]
.mul(&vars[1])
.add(&vars[2].exp())
.add(&Order2::constant(1.5));
let second_atom = vars[1].mul(&vars[3]).sub(&vars[0]);
let first_value = first_atom.value();
let second_exp = second_atom.value().exp();
let first_stack = [
first_value.ln(),
first_value.recip(),
-1.0 / (first_value * first_value),
0.0,
0.0,
];
let second_stack = [second_exp, second_exp, second_exp, second_exp, second_exp];
let dense = first_atom
.compose_unary(first_stack)
.add(&second_atom.compose_unary(second_stack));
let terms = [
Term {
first: first_stack[1],
second: first_stack[2],
gradient: *first_atom.g(),
hessian: *first_atom.h(),
},
Term {
first: second_stack[1],
second: second_stack[2],
gradient: *second_atom.g(),
hessian: *second_atom.h(),
},
];
let (value, gradient, hessian) = DynamicOrder2Accumulator::from_composed_sum(
K,
first_stack[0] + second_stack[0],
&terms,
)
.into_channels();
close(value, dense.value(), "dynamic value");
for row in 0..K {
close(
gradient[row],
dense.g()[row],
&format!("dynamic gradient[{row}]"),
);
for column in 0..K {
close(
hessian[row * K + column],
dense.h()[row][column],
&format!("dynamic Hessian[{row},{column}]"),
);
}
}
}
#[derive(Clone, Copy, Debug)]
struct FullTwoPattern;
impl HessianPattern<2, 3> for FullTwoPattern {
const PAIRS: [(usize, usize); 3] = [(0, 0), (0, 1), (1, 1)];
const PAIR_BITS: [[u128; 2]; 2] = hessian_pair_bits(Self::PAIRS);
}
#[test]
fn patterned_order2_matches_dense_order2() {
type Sparse = PatternedOrder2<FullTwoPattern, 2, 3>;
let dense_vars: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
let sparse_vars: [Sparse; 2] = std::array::from_fn(|a| Sparse::variable(SEED[a], a));
let dense = row_expr(&dense_vars);
let sparse = row_expr(&sparse_vars);
close(sparse.value(), dense.value(), "patterned value");
for i in 0..2 {
close(sparse.g()[i], dense.g()[i], &format!("patterned grad[{i}]"));
for j in 0..2 {
close(
sparse.h()[i][j],
dense.h()[i][j],
&format!("patterned hess[{i}][{j}]"),
);
}
}
}
#[test]
fn compose_unary_with_scalar_seam_bit_identical() {
fn rand_unit(state: &mut u64) -> f64 {
let mut x = *state;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
*state = x;
2.0 * ((x >> 11) as f64 / ((1u64 << 53) as f64)) - 1.0
}
fn stack(u: f64) -> [f64; 5] {
[
u.sin(),
u.cos(),
(2.0 * u).sin(),
(0.5 * u).cos(),
u * u - 0.3,
]
}
fn run<const K: usize>(state: &mut u64, n: usize) -> usize {
for _ in 0..n {
let base = rand_unit(state);
let mut s = Order2::<K>::variable(base, 0);
for a in 1..K {
s = crate::nested_dual::JetField::mul(
&s,
&Order2::<K>::variable(rand_unit(state), a),
);
}
let with = s.compose_unary_with(stack);
let explicit = s.compose_unary(stack(s.value()));
assert_eq!(with.value().to_bits(), explicit.value().to_bits(), "value");
for a in 0..K {
assert_eq!(with.g()[a].to_bits(), explicit.g()[a].to_bits(), "g[{a}]");
for b in 0..K {
assert_eq!(
with.h()[a][b].to_bits(),
explicit.h()[a][b].to_bits(),
"h[{a}][{b}]"
);
}
}
}
n
}
let mut st = 0x9e37_79b9_7f4a_7c15u64;
let total = run::<2>(&mut st, 1100)
+ run::<3>(&mut st, 1100)
+ run::<4>(&mut st, 1100)
+ run::<9>(&mut st, 1100);
assert_eq!(total, 4400);
}
#[test]
fn one_seed_matches_tower_third_contracted() {
let t = tower();
let truth = t.third_contracted(&U);
let vars: [OneSeed<2>; 2] =
std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
let s = row_expr(&vars);
close(s.value(), t.v, "value");
for a in 0..2 {
for b in 0..2 {
close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
}
}
let third = s.contracted_third();
for a in 0..2 {
for b in 0..2 {
close(third[a][b], truth[a][b], &format!("third[{a}][{b}]"));
}
}
}
#[test]
fn fused_one_seed_channels_match_unfused_definition_932() {
const K: usize = 8;
fn random_scalar(state: &mut u64) -> f64 {
*state ^= *state << 13;
*state ^= *state >> 7;
*state ^= *state << 17;
((*state >> 11) as f64 / ((1_u64 << 53) as f64)) * 2.0 - 1.0
}
fn random_order2<const N: usize>(state: &mut u64) -> Order2<N> {
let mut tower = crate::jet_tower::Tower2::<N>::zero();
tower.v = random_scalar(state);
for axis in 0..N {
tower.g[axis] = random_scalar(state);
}
for row in 0..N {
for column in row..N {
let channel = random_scalar(state);
tower.h[row][column] = channel;
tower.h[column][row] = channel;
}
}
Order2(tower)
}
fn assert_channels_close<const N: usize>(
label: &str,
actual: &OneSeed<N>,
expected: &OneSeed<N>,
) {
for (part_label, actual_part, expected_part, require_exact_symmetry) in [
("base", &actual.base.0, &expected.base.0, false),
("eps", &actual.eps.0, &expected.eps.0, true),
] {
let check = |channel: &str, got: f64, want: f64| {
let tolerance = 2.0e-14 * got.abs().max(want.abs()).max(1.0);
assert!(
(got - want).abs() <= tolerance,
"{label} {part_label} {channel}: got={got:+.17e} want={want:+.17e}"
);
};
check("value", actual_part.v, expected_part.v);
for row in 0..N {
check(
&format!("gradient[{row}]"),
actual_part.g[row],
expected_part.g[row],
);
for column in 0..N {
check(
&format!("hessian[{row},{column}]"),
actual_part.h[row][column],
expected_part.h[row][column],
);
if require_exact_symmetry {
assert_eq!(
actual_part.h[row][column].to_bits(),
actual_part.h[column][row].to_bits(),
"{label} {part_label} Hessian symmetry at [{row},{column}]"
);
}
}
}
}
}
let mut state = 0x9320_1eed_5eed_cafe_u64;
for sample in 0..256 {
let left = OneSeed {
base: random_order2::<K>(&mut state),
eps: random_order2::<K>(&mut state),
};
let right = OneSeed {
base: random_order2::<K>(&mut state),
eps: random_order2::<K>(&mut state),
};
let fused_product = left.mul(&right);
let unfused_product = OneSeed {
base: left.base.mul(&right.base),
eps: left.base.mul(&right.eps).add(&left.eps.mul(&right.base)),
};
assert_channels_close(
&format!("sample {sample} product"),
&fused_product,
&unfused_product,
);
let derivatives: [f64; 5] = std::array::from_fn(|_| random_scalar(&mut state));
let fused_composition = left.compose_unary(derivatives);
let unfused_composition = OneSeed {
base: left.base.compose_unary(derivatives),
eps: left
.base
.compose_unary([
derivatives[1],
derivatives[2],
derivatives[3],
derivatives[4],
derivatives[4],
])
.mul(&left.eps),
};
assert_channels_close(
&format!("sample {sample} composition"),
&fused_composition,
&unfused_composition,
);
}
}
#[test]
fn two_seed_matches_tower_fourth_contracted() {
let t = tower();
let truth4 = t.fourth_contracted(&U, &V);
let truth3_u = t.third_contracted(&U);
let truth3_v = t.third_contracted(&V);
let vars: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
let s = row_expr(&vars);
close(s.value(), t.v, "value");
for a in 0..2 {
close(s.base.0.g[a], t.g[a], &format!("grad[{a}]"));
for b in 0..2 {
close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
close(
s.eps.h()[a][b],
truth3_u[a][b],
&format!("eps third_u[{a}][{b}]"),
);
close(
s.del.h()[a][b],
truth3_v[a][b],
&format!("del third_v[{a}][{b}]"),
);
}
}
let fourth = s.contracted_fourth();
for a in 0..2 {
for b in 0..2 {
close(fourth[a][b], truth4[a][b], &format!("fourth[{a}][{b}]"));
}
}
}
#[test]
fn generic_program_seam_matches_tower_for_every_channel() {
let t = tower();
let o2: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
let so2 = row_expr(&o2);
close(so2.value(), t.v, "seam order2 value");
let os: [OneSeed<2>; 2] =
std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
let third = row_expr(&os).contracted_third();
let truth3 = t.third_contracted(&U);
for a in 0..2 {
for b in 0..2 {
close(third[a][b], truth3[a][b], &format!("seam third[{a}][{b}]"));
}
}
let ts: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
let fourth = row_expr(&ts).contracted_fourth();
let truth4 = t.fourth_contracted(&U, &V);
for a in 0..2 {
for b in 0..2 {
close(
fourth[a][b],
truth4[a][b],
&format!("seam fourth[{a}][{b}]"),
);
}
}
}
#[test]
fn tower4_as_jetscalar_matches_program_tower_all_channels() {
let t = tower();
let vars: [Tower4<2>; 2] = std::array::from_fn(|a| Tower4::variable(SEED[a], a));
let s = row_expr(&vars);
close(s.v, t.v, "tower-jetscalar value");
for a in 0..2 {
close(s.g[a], t.g[a], &format!("tower-jetscalar grad[{a}]"));
for b in 0..2 {
close(
s.h[a][b],
t.h[a][b],
&format!("tower-jetscalar hess[{a}][{b}]"),
);
for c in 0..2 {
close(
s.t3[a][b][c],
t.t3[a][b][c],
&format!("tower-jetscalar t3[{a}][{b}][{c}]"),
);
for d in 0..2 {
close(
s.t4[a][b][c][d],
t.t4[a][b][c][d],
&format!("tower-jetscalar t4[{a}][{b}][{c}][{d}]"),
);
}
}
}
}
}
#[test]
fn runtime_directional_jets_match_fixed_packed_algebra_932() {
fn expression<'arena, S: RuntimeJetScalar<'arena>>(vars: &[S]) -> S {
let bilinear = vars[0].mul(&vars[1]);
let curved = vars[2].scale(0.7).add(&vars[3].mul(&vars[3]).scale(-0.2));
bilinear
.add(&curved)
.exp()
.mul(&vars[4].compose_unary([0.4, -0.3, 0.2, -0.1, 0.05]))
}
const K: usize = 5;
let values = [0.2, -0.7, 0.4, 1.1, -0.3];
let direction_u = [0.5, -0.2, 0.7, -0.4, 0.1];
let direction_v = [-0.3, 0.8, 0.2, 0.6, -0.5];
let close = |actual: f64, expected: f64| {
let tolerance = 1.0e-13 * (1.0 + actual.abs().max(expected.abs()));
assert!((actual - expected).abs() <= tolerance);
};
let fixed_one: Vec<FixedRuntimeJet<OneSeed<K>, K>> = (0..K)
.map(|axis| FixedRuntimeJet {
inner: OneSeed::seed_direction(values[axis], axis, direction_u[axis]),
})
.collect();
let arena_one = DynamicJetArena::new();
let dynamic_one: Vec<DynamicOneSeed<'_>> = (0..K)
.map(|axis| {
DynamicOneSeed::seed_direction(values[axis], axis, direction_u[axis], K, &arena_one)
})
.collect();
let fixed_third = expression(&fixed_one).into_inner().contracted_third();
let dynamic_third = expression(&dynamic_one);
for a in 0..K {
for b in 0..K {
assert_eq!(
dynamic_third.contracted_third()[a * K + b].to_bits(),
dynamic_third.contracted_third()[b * K + a].to_bits(),
"arena third Hessian must be exactly symmetric at ({a},{b})"
);
close(
dynamic_third.contracted_third()[a * K + b],
fixed_third[a][b],
);
}
}
let fixed_one_v: Vec<FixedRuntimeJet<OneSeed<K>, K>> = (0..K)
.map(|axis| FixedRuntimeJet {
inner: OneSeed::seed_direction(values[axis], axis, direction_v[axis]),
})
.collect();
let fixed_third_v = expression(&fixed_one_v).into_inner().contracted_third();
let batch_workspace = DynamicJetBatchWorkspace::new(2);
let directions = [direction_u, direction_v];
let batch_vars = batch_workspace.alloc_slice_fill_with(K, |axis| {
DynamicOneSeedBatch::seed_directions(values[axis], axis, K, &batch_workspace, |lane| {
directions[lane][axis]
})
});
let dynamic_batch = expression(batch_vars);
assert_eq!(dynamic_batch.lanes(), 2);
for lane in 0..2 {
let expected = if lane == 0 {
&fixed_third
} else {
&fixed_third_v
};
for a in 0..K {
for b in 0..K {
close(
dynamic_batch.contracted_third(lane)[a * K + b],
expected[a][b],
);
}
}
}
let fixed_two: Vec<FixedRuntimeJet<TwoSeed<K>, K>> = (0..K)
.map(|axis| FixedRuntimeJet {
inner: TwoSeed::seed(values[axis], axis, direction_u[axis], direction_v[axis]),
})
.collect();
let arena_two = DynamicJetArena::new();
let dynamic_two: Vec<DynamicTwoSeed<'_>> = (0..K)
.map(|axis| {
DynamicTwoSeed::seed(
values[axis],
axis,
direction_u[axis],
direction_v[axis],
K,
&arena_two,
)
})
.collect();
let fixed_fourth = expression(&fixed_two).into_inner().contracted_fourth();
let dynamic_fourth = expression(&dynamic_two);
for a in 0..K {
for b in 0..K {
close(
dynamic_fourth.contracted_fourth()[a * K + b],
fixed_fourth[a][b],
);
}
}
let fixed_two_swapped: Vec<FixedRuntimeJet<TwoSeed<K>, K>> = (0..K)
.map(|axis| {
FixedRuntimeJet::from_inner(TwoSeed::seed(
values[axis],
axis,
direction_v[axis],
direction_u[axis],
))
})
.collect();
let fixed_fourth_swapped = expression(&fixed_two_swapped)
.into_inner()
.contracted_fourth();
let pair_workspace = DynamicJetBatchWorkspace::new(2);
let direction_pairs = [(direction_u, direction_v), (direction_v, direction_u)];
let pair_vars = pair_workspace.alloc_slice_fill_with(K, |axis| {
DynamicTwoSeedBatch::seed_direction_pairs(
values[axis],
axis,
K,
&pair_workspace,
|lane| (direction_pairs[lane].0[axis], direction_pairs[lane].1[axis]),
)
});
let dynamic_pair_batch = expression(pair_vars);
assert_eq!(dynamic_pair_batch.lanes(), 2);
for lane in 0..2 {
let expected = if lane == 0 {
&fixed_fourth
} else {
&fixed_fourth_swapped
};
for a in 0..K {
for b in 0..K {
close(
dynamic_pair_batch.contracted_fourth(lane)[a * K + b],
expected[a][b],
);
}
}
}
}
#[test]
fn dynamic_jet_arena_compacts_fragmented_high_water_932() {
const WORDS_PER_ALLOCATION: usize = 1 << 17;
const ALLOCATIONS: usize = 6;
let mut arena = DynamicJetArena::new();
for lane in 0..ALLOCATIONS {
let allocation = arena.alloc_slice_fill_with(WORDS_PER_ALLOCATION, |_| lane as u64);
std::hint::black_box(allocation);
}
let fragmented_high_water = arena.allocated_bytes();
arena.reset();
let compact_high_water = arena.allocated_bytes();
assert!(
compact_high_water >= fragmented_high_water,
"compacted arena must retain the complete fragmented tape"
);
for lane in 0..ALLOCATIONS {
let allocation = arena.alloc_slice_fill_with(WORDS_PER_ALLOCATION, |_| lane as u64);
std::hint::black_box(allocation);
}
assert_eq!(
arena.allocated_bytes(),
compact_high_water,
"equal replay must fit in the compacted chunk"
);
arena.reset();
assert_eq!(
arena.allocated_bytes(),
compact_high_water,
"stable reset must retain the compacted chunk"
);
}
}
#[cfg(test)]
mod batch_tests {
use super::{
JetScalar, Lane, OneSeed, OneSeedBatch, OneSeedLane, Order2, Order2Batch, Order2Lane,
TwoSeed, TwoSeedBatch, TwoSeedLane,
};
use crate::nested_dual::JetField;
trait RowAlg<const K: usize>: Copy {
fn constant(c: f64) -> Self;
fn add(&self, o: &Self) -> Self;
fn sub(&self, o: &Self) -> Self;
fn mul(&self, o: &Self) -> Self;
fn scale(&self, s: f64) -> Self;
fn exp(&self) -> Self;
fn sqrt(&self) -> Self;
fn recip(&self) -> Self;
}
impl<const K: usize> RowAlg<K> for Order2<K> {
fn constant(c: f64) -> Self {
<Self as JetScalar<K>>::constant(c)
}
fn add(&self, o: &Self) -> Self {
crate::nested_dual::JetField::add(self, o)
}
fn sub(&self, o: &Self) -> Self {
crate::nested_dual::JetField::sub(self, o)
}
fn mul(&self, o: &Self) -> Self {
crate::nested_dual::JetField::mul(self, o)
}
fn scale(&self, s: f64) -> Self {
crate::nested_dual::JetField::scale(self, s)
}
fn exp(&self) -> Self {
JetScalar::exp(self)
}
fn sqrt(&self) -> Self {
JetScalar::sqrt(self)
}
fn recip(&self) -> Self {
JetScalar::recip(self)
}
}
impl<L: Lane, const K: usize> RowAlg<K> for Order2Lane<L, K> {
fn constant(c: f64) -> Self {
Order2Lane::constant(L::splat(c))
}
fn add(&self, o: &Self) -> Self {
Order2Lane::add(self, o)
}
fn sub(&self, o: &Self) -> Self {
Order2Lane::sub(self, o)
}
fn mul(&self, o: &Self) -> Self {
Order2Lane::mul(self, o)
}
fn scale(&self, s: f64) -> Self {
Order2Lane::scale(self, s)
}
fn exp(&self) -> Self {
Order2Lane::exp(self)
}
fn sqrt(&self) -> Self {
Order2Lane::sqrt(self)
}
fn recip(&self) -> Self {
Order2Lane::recip(self)
}
}
fn row_expr<const K: usize, A: RowAlg<K>>(p: &[A; K]) -> A {
let mut s = A::constant(0.3);
for a in 0..K {
let b = (a + 1) % K;
s = s.add(&p[a].mul(&p[b]).scale(0.1 + 0.05 * a as f64));
}
let e = s.exp();
let r = s.mul(&s).add(&A::constant(1.0)).sqrt();
let denom = e.add(&A::constant(2.0));
e.mul(&r).sub(&s.scale(0.5)).mul(&denom.recip())
}
fn rand_unit(state: &mut u64) -> f64 {
let mut x = *state;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
*state = x;
let u = (x >> 11) as f64 / ((1u64 << 53) as f64); 2.0 * u - 1.0
}
fn check_k<const K: usize>(state: &mut u64, batches: usize) -> usize {
let mut verified_rows = 0usize;
for _ in 0..batches {
let rows: [[f64; K]; 4] =
std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
let prod: [Order2<K>; 4] = std::array::from_fn(|r| {
let p: [Order2<K>; K] = std::array::from_fn(|a| Order2::variable(rows[r][a], a));
row_expr(&p)
});
let scal: [Order2Lane<f64, K>; 4] = std::array::from_fn(|r| {
let p: [Order2Lane<f64, K>; K] =
std::array::from_fn(|a| Order2Lane::variable(rows[r][a], a));
row_expr(&p)
});
let pbatch: [Order2Batch<K>; K] = std::array::from_fn(|a| {
let packed = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
Order2Batch::variable(packed, a)
});
let batch = row_expr(&pbatch);
for r in 0..4 {
let g = prod[r].0;
assert_eq!(scal[r].v.to_bits(), g.v.to_bits(), "K={K} scalar v");
let lr = batch.lane(r).0;
assert_eq!(lr.v.to_bits(), g.v.to_bits(), "K={K} batch lane {r} v");
for a in 0..K {
assert_eq!(
scal[r].g[a].to_bits(),
g.g[a].to_bits(),
"K={K} scalar g[{a}]"
);
assert_eq!(
lr.g[a].to_bits(),
g.g[a].to_bits(),
"K={K} batch lane {r} g[{a}]"
);
for b in 0..K {
assert_eq!(
scal[r].h[a][b].to_bits(),
g.h[a][b].to_bits(),
"K={K} scalar h[{a}][{b}]"
);
assert_eq!(
lr.h[a][b].to_bits(),
g.h[a][b].to_bits(),
"K={K} batch lane {r} h[{a}][{b}]"
);
}
}
verified_rows += 1;
}
}
verified_rows
}
#[test]
fn batch_lanes_bit_identical_to_scalar_per_row() {
let mut state = 0x9E37_79B9_7F4A_7C15_u64;
let mut verified = 0usize;
verified += check_k::<2>(&mut state, 2000);
verified += check_k::<3>(&mut state, 2000);
verified += check_k::<4>(&mut state, 2000);
verified += check_k::<9>(&mut state, 2000);
assert_eq!(verified, 4 * 2000 * 4, "every batch row must be verified");
}
impl<const K: usize> RowAlg<K> for OneSeed<K> {
fn constant(c: f64) -> Self {
<Self as JetScalar<K>>::constant(c)
}
fn add(&self, o: &Self) -> Self {
crate::nested_dual::JetField::add(self, o)
}
fn sub(&self, o: &Self) -> Self {
crate::nested_dual::JetField::sub(self, o)
}
fn mul(&self, o: &Self) -> Self {
crate::nested_dual::JetField::mul(self, o)
}
fn scale(&self, s: f64) -> Self {
crate::nested_dual::JetField::scale(self, s)
}
fn exp(&self) -> Self {
JetScalar::exp(self)
}
fn sqrt(&self) -> Self {
JetScalar::sqrt(self)
}
fn recip(&self) -> Self {
JetScalar::recip(self)
}
}
impl<L: Lane, const K: usize> RowAlg<K> for OneSeedLane<L, K> {
fn constant(c: f64) -> Self {
OneSeedLane::constant(L::splat(c))
}
fn add(&self, o: &Self) -> Self {
OneSeedLane::add(self, o)
}
fn sub(&self, o: &Self) -> Self {
OneSeedLane::sub(self, o)
}
fn mul(&self, o: &Self) -> Self {
OneSeedLane::mul(self, o)
}
fn scale(&self, s: f64) -> Self {
OneSeedLane::scale(self, s)
}
fn exp(&self) -> Self {
OneSeedLane::exp(self)
}
fn sqrt(&self) -> Self {
OneSeedLane::sqrt(self)
}
fn recip(&self) -> Self {
OneSeedLane::recip(self)
}
}
impl<const K: usize> RowAlg<K> for TwoSeed<K> {
fn constant(c: f64) -> Self {
<Self as JetScalar<K>>::constant(c)
}
fn add(&self, o: &Self) -> Self {
crate::nested_dual::JetField::add(self, o)
}
fn sub(&self, o: &Self) -> Self {
crate::nested_dual::JetField::sub(self, o)
}
fn mul(&self, o: &Self) -> Self {
crate::nested_dual::JetField::mul(self, o)
}
fn scale(&self, s: f64) -> Self {
crate::nested_dual::JetField::scale(self, s)
}
fn exp(&self) -> Self {
JetScalar::exp(self)
}
fn sqrt(&self) -> Self {
JetScalar::sqrt(self)
}
fn recip(&self) -> Self {
JetScalar::recip(self)
}
}
impl<L: Lane, const K: usize> RowAlg<K> for TwoSeedLane<L, K> {
fn constant(c: f64) -> Self {
TwoSeedLane::constant(L::splat(c))
}
fn add(&self, o: &Self) -> Self {
TwoSeedLane::add(self, o)
}
fn sub(&self, o: &Self) -> Self {
TwoSeedLane::sub(self, o)
}
fn mul(&self, o: &Self) -> Self {
TwoSeedLane::mul(self, o)
}
fn scale(&self, s: f64) -> Self {
TwoSeedLane::scale(self, s)
}
fn exp(&self) -> Self {
TwoSeedLane::exp(self)
}
fn sqrt(&self) -> Self {
TwoSeedLane::sqrt(self)
}
fn recip(&self) -> Self {
TwoSeedLane::recip(self)
}
}
fn check_oneseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
let mut rows_checked = 0;
for _ in 0..batches {
let rows: [[f64; K]; 4] =
std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
let u: [[f64; K]; 4] =
std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
let prod: [OneSeed<K>; 4] = std::array::from_fn(|r| {
let p: [OneSeed<K>; K] =
std::array::from_fn(|a| OneSeed::seed_direction(rows[r][a], a, u[r][a]));
row_expr(&p)
});
let scal: [OneSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
let p: [OneSeedLane<f64, K>; K] =
std::array::from_fn(|a| OneSeedLane::seed_direction(rows[r][a], a, u[r][a]));
row_expr(&p)
});
let pbatch: [OneSeedBatch<K>; K] = std::array::from_fn(|a| {
let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
OneSeedBatch::seed_direction(val, a, uu)
});
let batch = row_expr(&pbatch);
for r in 0..4 {
let want = prod[r].contracted_third();
let got_scal = scal[r].contracted_third();
let got_batch = batch.lane(r).contracted_third();
assert_eq!(
scal[r].base.v.to_bits(),
prod[r].base.value().to_bits(),
"OneSeed K={K} scalar value"
);
assert_eq!(
batch.lane(r).base.value().to_bits(),
prod[r].base.value().to_bits(),
"OneSeed K={K} batch lane {r} value"
);
for a in 0..K {
for b in 0..K {
assert_eq!(
got_scal[a][b].to_bits(),
want[a][b].to_bits(),
"OneSeed K={K} scalar third[{a}][{b}]"
);
assert_eq!(
got_batch[a][b].to_bits(),
want[a][b].to_bits(),
"OneSeed K={K} batch lane {r} third[{a}][{b}]"
);
}
}
rows_checked += 1;
}
}
rows_checked
}
fn check_twoseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
let mut rows_checked = 0;
for _ in 0..batches {
let rows: [[f64; K]; 4] =
std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
let u: [[f64; K]; 4] =
std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
let v: [[f64; K]; 4] =
std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
let prod: [TwoSeed<K>; 4] = std::array::from_fn(|r| {
let p: [TwoSeed<K>; K] =
std::array::from_fn(|a| TwoSeed::seed(rows[r][a], a, u[r][a], v[r][a]));
row_expr(&p)
});
let scal: [TwoSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
let p: [TwoSeedLane<f64, K>; K] =
std::array::from_fn(|a| TwoSeedLane::seed(rows[r][a], a, u[r][a], v[r][a]));
row_expr(&p)
});
let pbatch: [TwoSeedBatch<K>; K] = std::array::from_fn(|a| {
let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
let vv = wide::f64x4::new([v[0][a], v[1][a], v[2][a], v[3][a]]);
TwoSeedBatch::seed(val, a, uu, vv)
});
let batch = row_expr(&pbatch);
for r in 0..4 {
let want = prod[r].contracted_fourth();
let got_scal = scal[r].contracted_fourth();
let got_batch = batch.lane(r).contracted_fourth();
assert_eq!(
scal[r].base.v.to_bits(),
prod[r].base.value().to_bits(),
"TwoSeed K={K} scalar value"
);
assert_eq!(
batch.lane(r).base.value().to_bits(),
prod[r].base.value().to_bits(),
"TwoSeed K={K} batch lane {r} value"
);
for a in 0..K {
for b in 0..K {
assert_eq!(
got_scal[a][b].to_bits(),
want[a][b].to_bits(),
"TwoSeed K={K} scalar fourth[{a}][{b}]"
);
assert_eq!(
got_batch[a][b].to_bits(),
want[a][b].to_bits(),
"TwoSeed K={K} batch lane {r} fourth[{a}][{b}]"
);
}
}
rows_checked += 1;
}
}
rows_checked
}
#[test]
fn oneseed_lanes_contracted_third_bit_identical() {
let mut state = 0x1234_5678_9ABC_DEF0_u64;
let batches = 2000;
let rows_checked = check_oneseed::<2>(&mut state, batches)
+ check_oneseed::<3>(&mut state, batches)
+ check_oneseed::<4>(&mut state, batches)
+ check_oneseed::<9>(&mut state, batches);
assert_eq!(rows_checked, 4 * batches * 4);
}
#[test]
fn twoseed_lanes_contracted_fourth_bit_identical() {
let mut state = 0x0FED_CBA9_8765_4321_u64;
let batches = 2000;
let rows_checked = check_twoseed::<2>(&mut state, batches)
+ check_twoseed::<3>(&mut state, batches)
+ check_twoseed::<4>(&mut state, batches)
+ check_twoseed::<9>(&mut state, batches);
assert_eq!(rows_checked, 4 * batches * 4);
}
}
#[cfg(test)]
mod unit_tests {
use super::{
DynamicJetArena, DynamicOrder2, JetScalar, OneSeed, Order1, Order2, RuntimeJetScalar,
filtered_implicit_solve_scalar,
};
use crate::nested_dual::{Dual2, JetField};
fn family_program<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
let xy = x.mul(y);
let exponential = theta.mul(&xy).exp();
let theta_squared_x_squared = theta.mul(theta).mul(&x.mul(x)).scale(0.375);
let theta_y_cubed = theta.mul(&y.mul(y).mul(y)).scale(-0.2);
exponential
.add(&theta_squared_x_squared)
.add(&theta_y_cubed)
}
fn analytic_family_first<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
let xy = x.mul(y);
let exponential = theta.mul(&xy).exp();
xy.mul(&exponential)
.add(&theta.mul(&x.mul(x)).scale(0.75))
.add(&y.mul(y).mul(y).scale(-0.2))
}
fn analytic_family_second<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
let xy = x.mul(y);
let exponential = theta.mul(&xy).exp();
xy.mul(&xy).mul(&exponential).add(&x.mul(x).scale(0.75))
}
fn assert_channel_close(actual: f64, expected: f64, channel: &str) {
let tolerance = 256.0 * f64::EPSILON * (1.0 + actual.abs().max(expected.abs()));
assert!(
(actual - expected).abs() <= tolerance,
"{channel}: actual={actual:.17e}, expected={expected:.17e}, tolerance={tolerance:.3e}"
);
}
fn assert_order2_channels<const K: usize>(
actual: &Order2<K>,
expected: &Order2<K>,
prefix: &str,
) {
assert_channel_close(actual.value(), expected.value(), &format!("{prefix}.value"));
for a in 0..K {
assert_channel_close(actual.g()[a], expected.g()[a], &format!("{prefix}.g[{a}]"));
for b in 0..K {
assert_channel_close(
actual.h()[a][b],
expected.h()[a][b],
&format!("{prefix}.h[{a}][{b}]"),
);
}
}
}
#[test]
fn runtime_shaped_value_primitives_are_exact_and_skip_constant_composition_932() {
use std::time::Instant;
const K: usize = 48;
let arena = DynamicJetArena::new();
let variable = DynamicOrder2::variable(0.75, 7, K, &arena);
let constant = variable.constant_like(1.25);
assert_eq!(constant.value().to_bits(), 1.25_f64.to_bits());
assert_eq!(constant.dimension(), K);
assert!(constant.g().iter().all(|&channel| channel == 0.0));
assert!(constant.h().iter().all(|&channel| channel == 0.0));
let replaced = variable.with_value(-2.5);
assert_eq!(replaced.value().to_bits(), (-2.5_f64).to_bits());
assert_eq!(replaced.g(), variable.g());
assert_eq!(replaced.h(), variable.h());
fn best_ns(mut evaluate: impl FnMut(f64) -> f64, iterations: usize) -> f64 {
let mut best = f64::INFINITY;
for _ in 0..5 {
let mut checksum = 0.0_f64;
let started = Instant::now();
for _ in 0..iterations {
checksum += evaluate(0.75 + checksum * 1e-18);
}
assert!(checksum.is_finite());
best = best.min(started.elapsed().as_secs_f64());
}
best * 1e9 / iterations as f64
}
let iterations = if cfg!(debug_assertions) { 200 } else { 20_000 };
let mut direct_arena = DynamicJetArena::new();
let direct_ns = best_ns(
|value| {
direct_arena.reset();
let variable = DynamicOrder2::variable(value, 7, K, &direct_arena);
let constant = variable.constant_like(1.25);
constant.value() + constant.g()[K - 1] + constant.h()[K * K - 1]
},
iterations,
);
let mut composed_arena = DynamicJetArena::new();
let composed_ns = best_ns(
|value| {
composed_arena.reset();
let variable = DynamicOrder2::variable(value, 7, K, &composed_arena);
let constant = variable.compose_unary([1.25, 0.0, 0.0, 0.0, 0.0]);
constant.value() + constant.g()[K - 1] + constant.h()[K * K - 1]
},
iterations,
);
eprintln!(
"RUNTIME-CONSTANT-932 dimension={K} direct={direct_ns:.2} ns \
composed={composed_ns:.2} ns composed_over_direct={:.6}",
composed_ns / direct_ns,
);
}
#[test]
fn dual2_order2_extracts_exact_family_value_gradient_hessian_channels() {
const K: usize = 2;
let x0 = 0.7;
let y0 = -0.45;
let theta0 = 0.6;
let x = <Dual2<Order2<K>> as JetScalar<K>>::variable(x0, 0);
let y = <Dual2<Order2<K>> as JetScalar<K>>::variable(y0, 1);
let theta = Dual2 {
v: Order2::constant(theta0),
g: Order2::constant(1.0),
h: Order2::constant(0.0),
};
let actual = family_program(&x, &y, &theta);
let reference_x = Order2::variable(x0, 0);
let reference_y = Order2::variable(y0, 1);
let reference_theta = Order2::constant(theta0);
let expected_first = analytic_family_first(&reference_x, &reference_y, &reference_theta);
let expected_second = analytic_family_second(&reference_x, &reference_y, &reference_theta);
assert_order2_channels(&actual.g, &expected_first, "family_first");
assert_order2_channels(&actual.h, &expected_second, "family_second");
}
#[test]
fn dual2_oneseed_extracts_exact_family_hessian_drift() {
const K: usize = 2;
let x0 = 0.7;
let y0 = -0.45;
let theta0 = 0.6;
let direction = [0.3, -0.8];
let mut x = <Dual2<OneSeed<K>> as JetScalar<K>>::variable(x0, 0);
let mut y = <Dual2<OneSeed<K>> as JetScalar<K>>::variable(y0, 1);
x.v.eps = Order2::constant(direction[0]);
y.v.eps = Order2::constant(direction[1]);
let theta = Dual2 {
v: OneSeed::constant(theta0),
g: OneSeed::constant(1.0),
h: OneSeed::constant(0.0),
};
let actual = family_program(&x, &y, &theta);
let reference_x = OneSeed::seed_direction(x0, 0, direction[0]);
let reference_y = OneSeed::seed_direction(y0, 1, direction[1]);
let reference_theta = OneSeed::constant(theta0);
let expected = analytic_family_first(&reference_x, &reference_y, &reference_theta);
assert_order2_channels(&actual.g.eps, &expected.eps, "family_first_drift");
}
#[test]
fn order2_constant_has_zero_derivatives() {
let s = Order2::<3>::constant(7.5);
assert_eq!(s.value(), 7.5);
for a in 0..3 {
assert_eq!(s.g()[a], 0.0, "grad[{a}] should be zero");
for b in 0..3 {
assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
}
}
}
#[test]
fn order2_variable_has_unit_gradient_in_seeded_slot() {
let x = -2.5_f64;
let s = Order2::<4>::variable(x, 2);
assert_eq!(s.value(), x);
for a in 0..4 {
let expected_g = if a == 2 { 1.0 } else { 0.0 };
assert_eq!(s.g()[a], expected_g, "grad[{a}]");
for b in 0..4 {
assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
}
}
}
#[test]
fn order2_add_sub_roundtrip() {
let p = Order2::<2>::variable(3.0, 0);
let q = Order2::<2>::variable(2.0, 1);
let pq = crate::nested_dual::JetField::add(&p, &q);
assert_eq!(pq.value(), 5.0, "add value");
let back = crate::nested_dual::JetField::sub(&pq, &q);
for a in 0..2 {
assert_eq!(back.g()[a], p.g()[a], "grad[{a}] roundtrip");
}
}
#[test]
fn order2_mul_satisfies_leibniz_rule() {
let pv = 3.0_f64;
let qv = -2.0_f64;
let p = Order2::<2>::variable(pv, 0);
let q = Order2::<2>::variable(qv, 1);
let pq = crate::nested_dual::JetField::mul(&p, &q);
assert_eq!(pq.value(), pv * qv, "value = p·q");
assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
assert_eq!(pq.h()[0][1], 1.0, "∂²(p·q)/∂p∂q = 1");
assert_eq!(pq.h()[1][0], 1.0, "∂²(p·q)/∂q∂p = 1 (symmetric)");
assert_eq!(pq.h()[0][0], 0.0, "∂²(p·q)/∂p² = 0");
assert_eq!(pq.h()[1][1], 0.0, "∂²(p·q)/∂q² = 0");
}
#[test]
fn order2_scale_multiplies_all_channels() {
let p = Order2::<2>::variable(4.0, 0);
let s = 2.5_f64;
let ps = crate::nested_dual::JetField::scale(&p, s);
assert_eq!(ps.value(), 4.0 * s);
assert_eq!(ps.g()[0], 1.0 * s);
assert_eq!(ps.g()[1], 0.0);
}
#[test]
fn order2_exp_derivative_stack_correct() {
let p0 = 1.0_f64;
let p = Order2::<1>::variable(p0, 0);
let ep = JetScalar::exp(&p);
let e = p0.exp();
assert!((ep.value() - e).abs() < 1e-15, "exp value");
assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p) = exp(p)");
assert!((ep.h()[0][0] - e).abs() < 1e-15, "d²/dp² exp(p) = exp(p)");
}
#[test]
fn order2_ln_derivative_stack_correct() {
let p0 = 2.0_f64;
let p = Order2::<1>::variable(p0, 0);
let lnp = JetScalar::ln(&p);
assert!((lnp.value() - p0.ln()).abs() < 1e-15, "ln value");
assert!((lnp.g()[0] - 1.0 / p0).abs() < 1e-15, "d/dp ln(p) = 1/p");
assert!(
(lnp.h()[0][0] - (-1.0 / (p0 * p0))).abs() < 1e-15,
"d²/dp² ln(p) = -1/p²"
);
}
#[test]
fn dynamic_order2_ln_uses_runtime_scalar_derivative_stack() {
let p0 = 2.0_f64;
let arena = DynamicJetArena::new();
let p = DynamicOrder2::variable(p0, 0, 1, &arena);
let lnp = RuntimeJetScalar::ln(&p);
assert!((lnp.value() - p0.ln()).abs() < 1e-15, "ln value");
assert!((lnp.g()[0] - 1.0 / p0).abs() < 1e-15, "d/dp ln(p) = 1/p");
assert!(
(lnp.h_at(0, 0) - (-1.0 / (p0 * p0))).abs() < 1e-15,
"d²/dp² ln(p) = -1/p²"
);
}
#[test]
fn order2_exp_ln_roundtrip_at_value() {
let p0 = 0.8_f64;
let p = Order2::<1>::variable(p0, 0);
let roundtrip = JetScalar::ln(&JetScalar::exp(&p));
assert!((roundtrip.value() - p0).abs() < 1e-14, "ln(exp(p)) ≈ p");
}
#[test]
fn order1_constant_has_zero_gradient() {
let s = Order1::<3>::constant(-5.0);
assert_eq!(s.value(), -5.0);
for a in 0..3 {
assert_eq!(s.g()[a], 0.0, "g[{a}] should be zero");
}
}
#[test]
fn order1_variable_has_unit_gradient_in_seeded_slot() {
let s = Order1::<3>::variable(2.0, 1);
assert_eq!(s.value(), 2.0);
assert_eq!(s.g()[0], 0.0);
assert_eq!(s.g()[1], 1.0);
assert_eq!(s.g()[2], 0.0);
}
#[test]
fn order1_mul_satisfies_product_rule() {
let pv = 3.0_f64;
let qv = -2.0_f64;
let p = Order1::<2>::variable(pv, 0);
let q = Order1::<2>::variable(qv, 1);
let pq = crate::nested_dual::JetField::mul(&p, &q);
assert_eq!(pq.value(), pv * qv);
assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
}
#[test]
fn order1_exp_has_correct_value_and_gradient() {
let p0 = 0.5_f64;
let p = Order1::<2>::variable(p0, 0);
let ep = JetScalar::exp(&p);
let e = p0.exp();
assert!((ep.value() - e).abs() < 1e-15, "exp value");
assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p)");
assert_eq!(ep.g()[1], 0.0, "irrelevant gradient slot is zero");
}
#[test]
fn order1_and_order2_agree_on_value_and_gradient() {
let p0 = 1.3_f64;
let q0 = -0.7_f64;
let p1 = Order1::<2>::variable(p0, 0);
let q1 = Order1::<2>::variable(q0, 1);
let expr1 = JetScalar::exp(&crate::nested_dual::JetField::add(
&crate::nested_dual::JetField::mul(&p1, &q1),
&p1,
));
let p2 = Order2::<2>::variable(p0, 0);
let q2 = Order2::<2>::variable(q0, 1);
let expr2 = JetScalar::exp(&crate::nested_dual::JetField::add(
&crate::nested_dual::JetField::mul(&p2, &q2),
&p2,
));
assert!(
(expr1.value() - expr2.value()).abs() < 1e-14,
"value mismatch"
);
for a in 0..2 {
assert!(
(expr1.g()[a] - expr2.g()[a]).abs() < 1e-14,
"gradient[{a}] mismatch"
);
}
}
#[test]
fn filtered_implicit_solve_linear_constraint_gives_exact_jet() {
let theta0 = 3.0_f64;
let theta = Order2::<1>::variable(theta0, 0);
let a = filtered_implicit_solve_scalar::<1, Order2<1>>(theta0, 1.0, 2, |a_jet| {
crate::nested_dual::JetField::sub(a_jet, &theta)
});
assert!((a.value() - theta0).abs() < 1e-14, "value = theta0");
assert!((a.g()[0] - 1.0).abs() < 1e-14, "gradient = 1");
assert!(a.h()[0][0].abs() < 1e-14, "hessian = 0");
}
#[test]
fn filtered_implicit_solve_quadratic_constraint_matches_analytic_derivatives() {
let theta0 = 4.0_f64;
let a0 = theta0.sqrt();
let inv_fa = 1.0 / (2.0 * a0);
let theta = Order2::<1>::variable(theta0, 0);
let a = filtered_implicit_solve_scalar::<1, Order2<1>>(a0, inv_fa, 2, |a_jet| {
let aa = crate::nested_dual::JetField::mul(a_jet, a_jet);
crate::nested_dual::JetField::sub(&aa, &theta)
});
let tol = 1e-12;
assert!((a.value() - a0).abs() < tol, "value = sqrt(theta0)");
let expected_g = 0.5 / a0;
assert!(
(a.g()[0] - expected_g).abs() < tol,
"da/dtheta = 1/(2*sqrt)"
);
let expected_h = -0.25 / (theta0 * a0);
assert!(
(a.h()[0][0] - expected_h).abs() < tol,
"d2a/dtheta2 = -1/(4*theta^1.5)"
);
}
}