tenferro-internal-ops 0.8.0

//! Unit tests for elementwise AD helper rules.

use computegraph::graph::GraphBuilder;
use computegraph::types::{OperationRole, ValueKey, ValueRef};
use tenferro_tensor::{CompareDir, DType};

use crate::ad::context::{ShapeGuardContext, TensorMeta};
use crate::input_key::TensorInputKey;
use crate::std_tensor_op::StdTensorOp;
use crate::SymDim;

fn tensor_input(id: u64) -> TensorInputKey {
    TensorInputKey::User { id }
}

fn input_key(id: u64) -> ValueKey<StdTensorOp> {
    ValueKey::Input(tensor_input(id))
}

#[test]
fn zero_like_covers_scalar_and_broadcasted_dtype_paths() {
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let anchor = builder.add_input(tensor_input(1));

    let dtype_cases = [
        (DType::F32, 4),
        (DType::F64, 8),
        (DType::I64, 8),
        (DType::C32, 8),
        (DType::C64, 16),
    ];
    let mut scalar_zeros = Vec::new();
    for (dtype, byte_len) in dtype_cases {
        let zero =
            super::super::zeros::build_zero_like(&mut builder, dtype, ValueRef::Local(anchor), 0);
        scalar_zeros.push((zero, dtype, byte_len));
    }

    let vector_zero =
        super::super::zeros::build_zero_like(&mut builder, DType::F64, ValueRef::Local(anchor), 2);
    let graph = builder.build();

    for (zero, dtype, byte_len) in scalar_zeros {
        let (op_id, _) = graph.values()[zero]
            .producer
            .expect("zero scalar must be produced by a constant op");
        let op = &graph.operations()[op_id];
        match &op.operation {
            StdTensorOp::Constant { dtype: got, bytes } => {
                assert_eq!(*got, dtype);
                assert_eq!(bytes.len(), byte_len);
                assert!(bytes.iter().all(|byte| *byte == 0));
            }
            other => panic!("expected constant zero, got {other:?}"),
        }
    }

    let (broadcast, _) = graph.values()[vector_zero]
        .producer
        .expect("ranked zero must be produced by broadcast");
    let op = &graph.operations()[broadcast];
    assert_eq!(
        op.operation,
        StdTensorOp::BroadcastInDim {
            shape: vec![
                crate::dim_expr::DimExpr::InputDim {
                    input_idx: 1,
                    axis: 0
                },
                crate::dim_expr::DimExpr::InputDim {
                    input_idx: 1,
                    axis: 1
                },
            ],
            dims: vec![],
        }
    );
    assert_eq!(op.inputs[1], ValueRef::Local(anchor));
}

#[test]
fn symbolic_zero_carries_aval_until_instantiated() {
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let anchor = builder.add_input(tensor_input(1));
    let meta = TensorMeta::exact(DType::C64, vec![SymDim::from(2usize), SymDim::from(3usize)]);

    let zero = super::super::zeros::SymbolicZero::from_meta(ValueRef::Local(anchor), &meta);
    assert_eq!(zero.dtype(), DType::C64);
    assert_eq!(zero.rank(), 2);
    assert_eq!(zero.anchor(), &ValueRef::Local(anchor));
    assert!(builder.build().operations().is_empty());

    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let anchor = builder.add_input(tensor_input(2));
    let zero = super::super::zeros::SymbolicZero::from_meta(ValueRef::Local(anchor), &meta);
    let materialized = zero.instantiate(&mut builder);
    let graph = builder.build();

    let (op_id, _) = graph.values()[materialized]
        .producer
        .expect("ranked symbolic zero should materialize as broadcast");
    assert!(matches!(
        graph.operations()[op_id].operation,
        StdTensorOp::BroadcastInDim { .. }
    ));
}

#[test]
fn one_like_uses_dtype_aware_constant_without_exp_shortcut() {
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let anchor = builder.add_input(tensor_input(1));

    let scalar_one =
        super::super::zeros::build_one_like(&mut builder, DType::I64, ValueRef::Local(anchor), 0);
    let vector_one =
        super::super::zeros::build_one_like(&mut builder, DType::F64, ValueRef::Local(anchor), 2);
    let graph = builder.build();

    assert!(
        graph
            .operations()
            .iter()
            .all(|op| op.operation != StdTensorOp::Exp),
        "one-like seed helper must not synthesize constants through Exp"
    );

    let (scalar_op_id, _) = graph.values()[scalar_one]
        .producer
        .expect("scalar one must be produced by a constant op");
    match graph.operations()[scalar_op_id].operation.clone() {
        StdTensorOp::Constant { dtype, bytes } => {
            assert_eq!(dtype, DType::I64);
            assert_eq!(bytes, 1_i64.to_le_bytes().to_vec());
        }
        other => panic!("expected constant one, got {other:?}"),
    }

    let (broadcast_op_id, _) = graph.values()[vector_one]
        .producer
        .expect("ranked one must be produced by broadcast");
    let op = &graph.operations()[broadcast_op_id];
    assert_eq!(
        op.operation,
        StdTensorOp::BroadcastInDim {
            shape: vec![
                crate::dim_expr::DimExpr::InputDim {
                    input_idx: 1,
                    axis: 0
                },
                crate::dim_expr::DimExpr::InputDim {
                    input_idx: 1,
                    axis: 1
                },
            ],
            dims: vec![],
        }
    );
    assert_eq!(op.inputs[1], ValueRef::Local(anchor));
}

#[test]
fn linearize_elementwise_inactive_inputs_return_none_without_ops() {
    let mut ctx = ShapeGuardContext::default();
    let keys = vec![input_key(1), input_key(2), input_key(3)];

    for op in [
        StdTensorOp::Div,
        StdTensorOp::Abs,
        StdTensorOp::Sign,
        StdTensorOp::Maximum,
        StdTensorOp::Minimum,
        StdTensorOp::Select,
        StdTensorOp::Clamp,
    ] {
        let mut builder = GraphBuilder::<StdTensorOp>::new();
        let result = op
            .jvp_rule(
                &mut builder,
                &keys,
                &[input_key(4)],
                &[None, None, None],
                &mut ctx,
            )
            .unwrap();
        assert_eq!(result, vec![None], "{op:?}");
        assert!(
            builder.build().operations().is_empty(),
            "inactive {op:?} should not emit linearized ops"
        );
    }
}

#[test]
fn linearize_div_with_two_active_inputs_sums_both_terms() {
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let dx = builder.add_input(tensor_input(10));
    let dy = builder.add_input(tensor_input(11));
    let mut ctx = ShapeGuardContext::default();
    let op = StdTensorOp::Div;

    let result = op
        .jvp_rule(
            &mut builder,
            &[input_key(12), input_key(13)],
            &[input_key(14)],
            &[Some(dx), Some(dy)],
            &mut ctx,
        )
        .unwrap();

    assert!(result[0].is_some());
    let graph = builder.build();
    assert_eq!(
        graph.operations().last().map(|op| &op.operation),
        Some(&StdTensorOp::Add)
    );
}

#[test]
fn linearize_extrema_uses_jax_balanced_tie_masks() {
    let mut ctx = ShapeGuardContext::default();

    let mut max_builder = GraphBuilder::<StdTensorOp>::new();
    let dx = max_builder.add_input(tensor_input(20));
    let result = StdTensorOp::Maximum
        .jvp_rule(
            &mut max_builder,
            &[input_key(21), input_key(22)],
            &[],
            &[Some(dx), None],
            &mut ctx,
        )
        .unwrap();
    assert!(result[0].is_some());
    let graph = max_builder.build();
    assert_eq!(graph.operations()[0].operation, StdTensorOp::Maximum);
    assert_eq!(
        graph
            .operations()
            .iter()
            .filter(|op| op.operation == StdTensorOp::Compare(CompareDir::Eq))
            .count(),
        2
    );
    assert!(
        graph
            .operations()
            .iter()
            .filter(|op| op.operation == StdTensorOp::Select)
            .count()
            >= 2
    );
    assert!(graph
        .operations()
        .iter()
        .any(|op| op.operation == StdTensorOp::Div));

    let mut min_builder = GraphBuilder::<StdTensorOp>::new();
    let dy = min_builder.add_input(tensor_input(23));
    let result = StdTensorOp::Minimum
        .jvp_rule(
            &mut min_builder,
            &[input_key(24), input_key(25)],
            &[],
            &[None, Some(dy)],
            &mut ctx,
        )
        .unwrap();
    assert!(result[0].is_some());
    let graph = min_builder.build();
    assert_eq!(graph.operations()[0].operation, StdTensorOp::Minimum);
    assert_eq!(
        graph
            .operations()
            .iter()
            .filter(|op| op.operation == StdTensorOp::Compare(CompareDir::Eq))
            .count(),
        2
    );
    assert!(
        graph
            .operations()
            .iter()
            .filter(|op| op.operation == StdTensorOp::Select)
            .count()
            >= 2
    );
    assert!(graph
        .operations()
        .iter()
        .any(|op| op.operation == StdTensorOp::Div));
}

#[test]
fn linearize_extrema_mask_zero_is_a_dtype_constant_not_a_negated_value() {
    // A non-finite active value (for example an `Inf` tangent) must not be used to
    // synthesize the mask zero: `x + (-x)` evaluates to `NaN` there. The mask
    // seed must be a true dtype-aware `Constant` zero instead.
    let mut ctx = ShapeGuardContext::default();

    for (op, active_id, lhs_id, rhs_id) in [
        (StdTensorOp::Maximum, 60_u64, 61_u64, 62_u64),
        (StdTensorOp::Minimum, 63, 64, 65),
    ] {
        let mut builder = GraphBuilder::<StdTensorOp>::new();
        let active = builder.add_input(tensor_input(active_id));
        let result = op
            .jvp_rule(
                &mut builder,
                &[input_key(lhs_id), input_key(rhs_id)],
                &[],
                &[Some(active), None],
                &mut ctx,
            )
            .unwrap();
        assert!(result[0].is_some());
        let graph = builder.build();
        assert!(
            !graph
                .operations()
                .iter()
                .any(|op| op.operation == StdTensorOp::Neg),
            "the mask zero must be a dtype constant, not a negated active value"
        );
        assert!(
            graph.operations().iter().any(|op| matches!(
                &op.operation,
                StdTensorOp::Constant { dtype: DType::F64, bytes }
                    if bytes.iter().all(|byte| *byte == 0)
            )),
            "the mask zero must be emitted as a zero `Constant`"
        );
    }
}

#[test]
fn linearize_clamp_builds_nested_selects_for_active_bounds() {
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let dx = builder.add_input(tensor_input(30));
    let dlower = builder.add_input(tensor_input(31));
    let dupper = builder.add_input(tensor_input(32));
    let mut ctx = ShapeGuardContext::default();

    let result = StdTensorOp::Clamp
        .jvp_rule(
            &mut builder,
            &[input_key(33), input_key(34), input_key(35)],
            &[],
            &[Some(dx), Some(dlower), Some(dupper)],
            &mut ctx,
        )
        .unwrap();

    assert!(result[0].is_some());
    let graph = builder.build();
    assert!(
        graph
            .operations()
            .iter()
            .filter(|op| op.operation == StdTensorOp::Select)
            .count()
            >= 2
    );
}

#[test]
fn transpose_elementwise_handles_missing_or_inactive_cotangents() {
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let mut ctx = ShapeGuardContext::default();
    let inputs = vec![
        ValueRef::External(input_key(40)),
        ValueRef::External(input_key(41)),
    ];

    assert_eq!(
        StdTensorOp::Div
            .transpose_rule(
                &mut builder,
                &[None],
                &inputs,
                &OperationRole::Primary,
                &mut ctx
            )
            .unwrap(),
        vec![None, None]
    );
    let abs_ct = builder.add_input(tensor_input(42));
    assert_eq!(
        StdTensorOp::Abs
            .transpose_rule(
                &mut builder,
                &[Some(abs_ct)],
                &[inputs[0].clone()],
                &OperationRole::Primary,
                &mut ctx,
            )
            .unwrap(),
        vec![None]
    );
    assert_eq!(
        StdTensorOp::Maximum
            .transpose_rule(
                &mut builder,
                &[None],
                &inputs,
                &OperationRole::Linearized {
                    active_mask: vec![true, true],
                },
                &mut ctx,
            )
            .unwrap(),
        vec![None, None]
    );
}

#[test]
fn transpose_select_splits_cotangent_only_to_active_value_inputs() {
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let ct = builder.add_input(tensor_input(50));
    let mut ctx = ShapeGuardContext::default();
    let inputs = vec![
        ValueRef::External(input_key(51)),
        ValueRef::External(input_key(52)),
        ValueRef::External(input_key(53)),
    ];

    let result = StdTensorOp::Select
        .transpose_rule(
            &mut builder,
            &[Some(ct)],
            &inputs,
            &OperationRole::Linearized {
                active_mask: vec![false, true, false],
            },
            &mut ctx,
        )
        .unwrap();

    assert_eq!(result[0], None);
    assert!(result[1].is_some());
    assert_eq!(result[2], None);
    let graph = builder.build();
    assert!(graph
        .operations()
        .iter()
        .any(|op| op.operation == StdTensorOp::Select));
}

#[test]
fn transpose_clamp_covers_lower_only_and_inner_paths() {
    let mut ctx = ShapeGuardContext::default();
    let inputs = vec![
        ValueRef::External(input_key(60)),
        ValueRef::External(input_key(61)),
        ValueRef::External(input_key(62)),
    ];

    let mut lower_builder = GraphBuilder::<StdTensorOp>::new();
    let lower_ct = lower_builder.add_input(tensor_input(63));
    let result = StdTensorOp::Clamp
        .transpose_rule(
            &mut lower_builder,
            &[Some(lower_ct)],
            &inputs,
            &OperationRole::Linearized {
                active_mask: vec![false, true, false],
            },
            &mut ctx,
        )
        .unwrap();
    assert_eq!(result[0], None);
    assert!(result[1].is_some());
    assert_eq!(result[2], None);

    let mut full_builder = GraphBuilder::<StdTensorOp>::new();
    let full_ct = full_builder.add_input(tensor_input(64));
    let result = StdTensorOp::Clamp
        .transpose_rule(
            &mut full_builder,
            &[Some(full_ct)],
            &inputs,
            &OperationRole::Linearized {
                active_mask: vec![true, true, true],
            },
            &mut ctx,
        )
        .unwrap();
    assert!(result.iter().all(Option::is_some));
}

#[test]
fn linearize_select_emits_a_zero_constant_for_a_missing_branch() {
    // A missing Select branch must be filled with a dtype-aware constant zero,
    // never with an arithmetic zero derived from the other branch.
    let mut builder = GraphBuilder::<StdTensorOp>::new();
    let dx = builder.add_input(tensor_input(70));
    let mut ctx = ShapeGuardContext::default();

    let result = StdTensorOp::Select
        .jvp_rule(
            &mut builder,
            &[input_key(71), input_key(72), input_key(73)],
            &[],
            &[None, Some(dx), None],
            &mut ctx,
        )
        .unwrap();

    assert!(result[0].is_some());
    let graph = builder.build();
    assert!(!graph
        .operations()
        .iter()
        .any(|op| op.operation == StdTensorOp::Neg));
    assert!(graph.operations().iter().any(|op| matches!(
        &op.operation,
        StdTensorOp::Constant { bytes, .. } if bytes.iter().all(|byte| *byte == 0)
    )));
}

#[test]
fn transpose_extrema_emit_a_zero_constant_for_the_masked_operand() {
    for (op, ct_id, lhs_id, rhs_id) in [
        (StdTensorOp::Maximum, 80_u64, 81_u64, 82_u64),
        (StdTensorOp::Minimum, 83, 84, 85),
    ] {
        let mut builder = GraphBuilder::<StdTensorOp>::new();
        let ct = builder.add_input(tensor_input(ct_id));
        let mut ctx = ShapeGuardContext::default();
        let inputs = vec![
            ValueRef::External(input_key(lhs_id)),
            ValueRef::External(input_key(rhs_id)),
        ];

        let result = op
            .transpose_rule(
                &mut builder,
                &[Some(ct)],
                &inputs,
                &OperationRole::Linearized {
                    active_mask: vec![true, false],
                },
                &mut ctx,
            )
            .unwrap();

        assert!(result[0].is_some());
        assert!(result[1].is_none());
        let graph = builder.build();
        assert!(graph.operations().iter().any(|op| matches!(
            &op.operation,
            StdTensorOp::Constant { bytes, .. } if bytes.iter().all(|byte| *byte == 0)
        )));
    }
}