moose 0.2.2

Encrypted learning and data processing framework
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

use crate::computation::*;
use crate::logical::TensorShape;

/// The pass replaces the HostShape on the logical level with its Shape<Host> counterpart.
pub fn deprecated_shape_support(comp: Computation) -> anyhow::Result<Computation> {
    let mut operations = comp.operations;
    for op in operations.iter_mut() {
        // Recognize ops that have logical level kernels and may see a shape of a wrong type.
        match op.kind {
            Operator::Shape(ShapeOp {
                sig:
                    Signature::Unary(UnarySignature {
                        arg0: Ty::Tensor(_),
                        ret: ref mut ret_ty,
                    }),
            }) if *ret_ty == Ty::HostShape => {
                *ret_ty = Ty::Shape(TensorShape::Host);
            }
            Operator::Ones(OnesOp {
                sig:
                    Signature::Unary(UnarySignature {
                        arg0: ref mut arg0_ty,
                        ret: Ty::Tensor(_),
                    }),
            }) if *arg0_ty == Ty::HostShape => {
                *arg0_ty = Ty::Shape(TensorShape::Host);
            }
            Operator::Slice(SliceOp {
                sig:
                    Signature::Unary(UnarySignature {
                        arg0: Ty::HostShape,
                        ret: Ty::HostShape,
                    }),
                ..
            }) => {
                *op.kind.sig_mut() =
                    Signature::unary(Ty::Shape(TensorShape::Host), Ty::Shape(TensorShape::Host));
            }
            _ => {}
        };
    }
    Ok(Computation { operations })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::textual::ToTextual;
    use std::convert::TryInto;

    #[test]
    fn test_host_shape_replace() -> std::result::Result<(), anyhow::Error> {
        let source = r#"
        shape_0 = Shape: (Tensor<Fixed128(24, 40)>) -> HostShape (decrypt_0) @Host(bob)
        "#;

        let comp = deprecated_shape_support(source.try_into()?)?.to_textual();
        // The computation should now contain the modified type information
        assert_eq!(
            comp,
            "shape_0 = Shape: (Tensor<Fixed128(24, 40)>) -> Shape<Host> (decrypt_0) @Host(bob)"
        );
        Ok(())
    }
}