tenferro-ad 0.3.0

Eager runtime, eager tensors, and traced AD extension traits for tenferro.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use tenferro_ops::dim_expr::DimExpr;
use tenferro_ops::ShapeExtent;
use tenferro_runtime::SliceConfig;

use super::*;

pub(super) fn linearize_dynamic_shape(
    builder: &mut SemanticProgramBuilder,
    op: &CoreSemanticOp,
    primal_inputs: &[ProgramValue],
    tangent: AdValue,
) -> Result<AdValue, SemanticAdTransformError> {
    let AdValue::Value(tangent) = tangent else {
        return Ok(AdValue::Absent);
    };
    Ok(AdValue::Value(
        builder.add_op(op.clone(), &[tangent, primal_inputs[1]])?[0],
    ))
}

pub(super) fn dynamic_shape_vjp(
    builder: &mut SemanticProgramBuilder,
    op: &CoreSemanticOp,
    primal_inputs: &[ProgramValue],
    cotangent: AdValue,
    active_inputs: &[bool],
) -> Result<Vec<AdValue>, SemanticAdTransformError> {
    let mut result = vec![AdValue::Absent; 2];
    if !active_inputs[0] {
        return Ok(result);
    }
    let AdValue::Value(cotangent) = cotangent else {
        return Ok(result);
    };
    let value = match op {
        CoreSemanticOp::DynamicTruncate { axis } => builder.add_op(
            CoreSemanticOp::PadToMatch { axis: *axis },
            &[cotangent, primal_inputs[0]],
        )?[0],
        CoreSemanticOp::PadToMatch { axis } => {
            transpose_pad_to_match(builder, cotangent, primal_inputs[0], *axis)?
        }
        _ => return Err(unsupported_core(SemanticTransformRole::Vjp, op)),
    };
    result[0] = AdValue::Value(value);
    Ok(result)
}

fn transpose_pad_to_match(
    builder: &mut SemanticProgramBuilder,
    cotangent: ProgramValue,
    input: ProgramValue,
    axis: usize,
) -> Result<ProgramValue, SemanticAdTransformError> {
    let metadata = builder.value_metadata(input)?;
    if axis >= metadata.shape().len() {
        return Err(SemanticAdTransformError::UnsupportedMetadata {
            role: SemanticTransformRole::Vjp,
            message: format!(
                "pad-to-match axis {axis} is outside input rank {}",
                metadata.shape().len()
            ),
        });
    }
    let concrete_shape = metadata
        .shape()
        .iter()
        .map(|extent| match extent {
            ShapeExtent::Exact(DimExpr::Const(value)) => Some(*value),
            _ => None,
        })
        .collect::<Option<Vec<_>>>();
    if let Some(limits) = concrete_shape {
        let rank = limits.len();
        return Ok(builder.add_op(
            CoreSemanticOp::Slice(SliceConfig {
                starts: vec![0; rank],
                limits,
                strides: vec![1; rank],
            }),
            &[cotangent],
        )?[0]);
    }
    let size = builder.add_op(CoreSemanticOp::ShapeOf { axis }, &[input])?[0];
    Ok(builder.add_op(CoreSemanticOp::DynamicTruncate { axis }, &[cotangent, size])?[0])
}