tf-encrypted / moose

Secure distributed dataflow framework for encrypted machine learning and data processing
Apache License 2.0
59 stars 16 forks source link

Type checking and inference #333

Open mortendahl opened 3 years ago

mortendahl commented 3 years ago

We need the ability to both type check a computation, but also perform inference for missing types (Ty::UnknownTy).

mortendahl commented 3 years ago

Re type inference, I initially thought typical unification would be enough, but the problem is that the selective nature of operators (ie only certain instantiations of them are allowed) seems to lead to a somewhat complex constraint language. On top of that, we'd need to somehow generalize of our types so that eg unify(RingTensor<T>, RingTensor<U>) can reduce to unify(T, U), which is currently not the case since RingTensors are atomic in the Moose type system (enum Ty).

If possible I'd like to build for high cohesion around operators, including their logic for type checking/inference and semantics. The goal should be that adding a new operator is relatively straight-forward and low-risk. A complex constraint language and solver might work against this.

As an alternative to unification for type inference, I'm hence wondering if we should instead investigate an approach based on the monotone framework, where each operator is expected to implement its own monotone increasing type kernel for "improving" type signatures. As an example, the type kernel tk for RingAddOp might perform the following mappings:

This means that the contract between operators and the type inference algorithm can be based on the existing Ty and simply states that type kernels must be monotone increasing functions.

One question that remains, is whether we can use the information given by the kernel macros to automatically derive these type kernels:

kernel! {
    RingAddOp,
    [
        (HostPlacement, (Ring64Tensor, Ring64Tensor) -> Ring64Tensor => Self::kernel),
        (HostPlacement, (Ring128Tensor, Ring128Tensor) -> Ring128Tensor => Self::kernel),
    ]
}
mortendahl commented 3 years ago

Note that for type inference we probably want information to flow both forward and backwards, meaning the forward-only approach in the spike (where kernels are Box<dyn Fn(X, Y) -> Z>) will probably not fit with type inference (where we might need closures Box<dyn Fn(X, Y, Z) -> ?>).