ijzer_lib 0.1.1

Library for IJzer. Provides tools for tensors, parsing syntax tree of the IJ language and transpiling it to rust code.
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

//! SVD operator.
//! 
//! The SVD operator takes in a tensor and produces 3 tensors.
//! 
//! Example:
//! ```ijzer
//! (u, s, vt) = svd X
//! ```
//! 
//! 
//! 
use super::{check_ok_needed_outputs, gather_operands, ParseNode, ParseNodeFunctional};

use crate::ast_node::{ASTContext, Node, TokenSlice};
use crate::operations::Operation;
use crate::syntax_error::SyntaxError;
use crate::tokens::Token;
use crate::types::{FunctionSignature, IJType};
use anyhow::Result;
use std::rc::Rc;

pub struct Svd;
impl ParseNode for Svd {
    fn next_node(
        _op: Token,
        slice: TokenSlice,
        context: &mut ASTContext,
    ) -> Result<(Rc<Node>, TokenSlice)> {
        let (operands, rest) = gather_operands(vec![vec![IJType::Tensor(None)]], slice, context)?;
        let number_type = operands[0]
            .output_type
            .extract_number_type()
            .unwrap_or_default();

        Ok((
            Rc::new(Node::new(
                Operation::Svd,
                vec![operands[0].output_type.clone()],
                IJType::Group(vec![
                    IJType::Tensor(number_type.clone()),
                    IJType::Tensor(number_type.clone()),
                    IJType::Tensor(number_type),
                ]),
                operands,
                context,
            )?),
            rest,
        ))
    }
}

impl ParseNodeFunctional for Svd {
    fn next_node_functional_impl(
        _op: Token,
        slice: TokenSlice,
        context: &mut ASTContext,
        needed_outputs: Option<&[IJType]>,
    ) -> Result<(Vec<Rc<Node>>, TokenSlice)> {
        let rest = slice.move_start(1)?;
        let output_type = IJType::Group(vec![
            IJType::Tensor(None),
            IJType::Tensor(None),
            IJType::Tensor(None),
        ]);
        if !check_ok_needed_outputs(needed_outputs, &output_type) {
            return Err(SyntaxError::RequiredOutputsDoNotMatchFunctionOutputs(
                format!("{:?}", needed_outputs),
                output_type.to_string(),
            )
            .into());
        }

        let node = Rc::new(Node::new(
            Operation::Svd,
            vec![],
            IJType::Function(FunctionSignature::new(
                vec![IJType::Tensor(None)],
                output_type.clone(),
            )),
            vec![],
            context,
        )?);

        Ok((vec![node], rest))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::parse_str_no_context;

    #[test]
    fn test_svd() -> Result<()> {
        let (node, _) = parse_str_no_context("svd [1,2]")?;
        assert_eq!(node.op, Operation::Svd);
        assert_eq!(
            node.output_type,
            IJType::Group(vec![
                IJType::Tensor(None),
                IJType::Tensor(None),
                IJType::Tensor(None),
            ])
        );

        let (node, _) = parse_str_no_context(".~svd [1,2]")?;
        assert_eq!(node.op, Operation::Apply);
        assert_eq!(
            node.output_type,
            IJType::Group(vec![
                IJType::Tensor(None),
                IJType::Tensor(None),
                IJType::Tensor(None)
            ])
        );

        let (node, _) = parse_str_no_context("(u,s,vt) = svd [1,2]")?;
        assert_eq!(node.op, Operation::Assign);
        let lhs = node.operands[0].clone();
        assert_eq!(lhs.op, Operation::Group);
        let rhs = node.operands[1].clone();
        assert_eq!(rhs.op, Operation::Svd);
        Ok(())
    }
}