tenferro-internal-ops 0.8.0

Graph op vocabulary and AD rule implementations for the tenferro workspace.
Documentation 
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use std::any::Any;
use std::hash::Hasher;
use std::sync::Arc;

use tenferro_core_ops::{all_primitive_descriptors, PrimitiveOpKind};
use tenferro_tensor::{CompareDir, DType};

use crate::ext_op::ExtensionOp;
use crate::std_tensor_op::StdTensorOp;
use crate::SymDim;

#[derive(Clone, Debug)]
struct CatalogExt;

impl ExtensionOp for CatalogExt {
    fn family_id(&self) -> &'static str {
        "catalog.ext.v1"
    }

    fn payload_hash(&self, _hasher: &mut dyn Hasher) {}

    fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
        other.as_any().downcast_ref::<CatalogExt>().is_some()
    }

    fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
        Arc::new(self.clone())
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn input_count(&self) -> usize {
        1
    }

    fn output_count(&self) -> usize {
        1
    }

    fn infer_output_meta(
        &self,
        ctx: &mut crate::ExtensionShapeContext<'_>,
    ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
        Ok(vec![(ctx.input_dtype(0)?, ctx.input_shape(0)?.to_vec())])
    }
}

#[test]
fn std_tensor_op_maps_core_ops_to_catalog_kinds() {
    assert_eq!(
        StdTensorOp::Add.primitive_kind(),
        Some(PrimitiveOpKind::Add)
    );
    assert_eq!(
        StdTensorOp::Compare(CompareDir::Lt).primitive_kind(),
        Some(PrimitiveOpKind::Compare)
    );
    assert_eq!(
        StdTensorOp::Convert {
            from: DType::F32,
            to: DType::F64,
        }
        .primitive_kind(),
        Some(PrimitiveOpKind::Convert)
    );
}

#[test]
fn extension_ops_do_not_claim_core_kind() {
    let op = StdTensorOp::Extension(Arc::new(CatalogExt));

    assert_eq!(op.primitive_kind(), None);
}

#[test]
fn every_catalog_descriptor_has_a_std_tensor_op_variant() {
    for descriptor in all_primitive_descriptors() {
        let op = StdTensorOp::sample_from_kind(descriptor.kind);
        assert_eq!(
            op.primitive_kind(),
            Some(descriptor.kind),
            "StdTensorOp variant for {:?} disagrees with descriptor",
            descriptor.kind
        );
    }
}

#[test]
fn extension_ad_api_uses_linearize_and_transpose_terminology() {
    let source = include_str!("../ext_op.rs");

    for forbidden in [
        "ExtensionChainRule",
        "FruleBuilder",
        "RRuleBuilder",
        "register_extension_chain_rule",
        "frule",
        "rrule",
    ] {
        assert!(
            !source.contains(forbidden),
            "extension AD API should not expose ChainRules-style `{forbidden}` terminology"
        );
    }
}