Wondering static typed tensor

[jerry73204]

Inspired by Tensor Considered Harmful, I'm wondering if it's possible to build compile-time checked tensor type. There prior works to make things possible.

• typenum for typed numbers. It's necessary for CUDA device ordinal and dimension size.
• frunk for HList. Each dimension component can be named by a unit struct. HList acts as a Vec of distinct types that help us build the type-level dimension.
• The last is const generics. Though it still unstable and panics the compiler. It would replace typenum in the long run.

I made a little demo to demonstrate this idea. May we push it further:

• Introduce tensor operations like expand, multiply, flatten and more into type semantics.
• Since tch is backed on libtorch, a thorough test is required to verify the type design.

[LaurentMazare]

The demo looks very appealing, nice work! I certainly think that such approach are well suited for strongly typed languages like rust (or ocaml) and would provide some additional compile-time safety compared to the python front-end.

A first step could be to include the type of the elements in the tensor type. I gave this a shot in this branch of the ocaml version of tch-rs but haven't finished it. The main culprit is that most of the torch or ocaml apis are automatically generated from an ops description file but as far as I know this file currently lacks some information about the ops input/output types (see this issue about some ongoing cleanup of this ops description file). An alternative could be to manually write this typed layer on top of the current unsafe layer, this would also be an opportunity to fix #33.

Named dimensions on top of this would make it even better. Hopefully it's possible to come up with a reasonably complete api. I would be happy to see your prototype evolve. If there are some tch-rs changes that would make things easier feel free to create additional issues/PR!

[jerry73204]

I add more examples about static dimension permutation and broadcasting. The design is mostly inspired by the frunk author's article.

Suppose we have to implement dimension logic into worlds of types and traits, I'm not sure if these code can be totally generated from automated manner. ATen API itself seems to be lack of such information. (You may see complex broadcasting code) So far, I prefer to make named tensor a wrapper over tch tensor, and plus several debug assertions.

I will turn it into a standalone project. I expect some work there, and would welcome contributions.

[zeroexcuses]

While I would love to see compile time dimension working, I see a number of issues here:

1. With the current state of rustc, typenum can really slow down compile time. I've played with a minimal implementation inspired by typenum, and even that was still slow.

2. Throw in a few couple concat + convolutions, and we may need the rustc type checker to solve things like (m + n)2 = n2 + m*2.

For example, is A has dim m, B has dim n, and we do:

C = dup(concat(A, B)) D = concat(dup(B), dup(A)) // edit: swapping A & B order

E = C + D

to evaluate E, we need to ensure that C & D have the same dimension, which requires compile time checking that (m+n)2 = n2 + m*2 -- and it's not immediately obvious to me how rustc can do that.

[jerry73204]

1. I agree. In the long run, typenum will be replaced by const generics. Const generics is still unstable and is known to crash rustc, and there still no proper replacement to typenum traits (such as typenum::Unsigned). We can wait for it to evolve, and we continue on our progress.

2. I don't fully get your point. Suppose you're saying A and B are typed numbers. With typenum, A, B, C and D will be reduced to a typed list of bits. We could play some trick verify if C and D are equal.

type C = Prod<Sum<A, B>, U2>;
type D = Sum<Prod<B, U2>, Prod<A, U2>>;

let _: D = C::new(); // See if it compiles

[zeroexcuses]

Re 2:

Sorry for the confusion. Let me try this. Suppose we're dealing with just 1-d tensors everywhere.

fn <M, N> concat(m: &Tensor<M>, n: &Tensor<N>) -> Tensor<M+N> {
...
}

fn <M> dup(m: &Tensor<M>) -> Tensor<M+M> {
  concat(m, m)
}

fn <M, N> foo(m: &Tensor<M>, n: &Tensor<N>) {
  let C: Tensor<(M+N)+(M+N)> = dup(concat(m, n));
  let D: Tensor<(M+M)+(N+N)> = concat(dup(m), dup(n));
  let E = C + D; // we should be able to do element wise addition
}

In my limited understanding of the rustc type checker, it's not clear to me it can infer that Tensor<(M+N)+(M+N)> and Tensor<(M+M) + (N+N)> are the same types.

[zeroexcuses]

To make matters more complicated, once we throw in convolutions / pooling, we'll need rustc to be able to solve algebraic equalities involving +, -, div by constant, floor, ceiling.

[jerry73204]

For simplicity, I would assume the generics M, N are pure integers. typenum already provides a set of type operators to solve this issue. Typed number are reduced typed bits in typenum. Hence, most ordinary algebraic rules hold and it's not hard to infer M+N = N+M. Let me show the concat example that involves in typed number addition.

On the other hand, if some dimension size is only known in runtime, things would be more complicated. It happens when users would dynamically tweak the batch size. We would introduce variables (or placeholders), and solve rings of polynomials. I'm currently looking at frunk to see further inspirations.

I move most of my work in this repo. It's still in alpha and I'm glad for contributions. https://github.com/jerry73204/tch-typed-tensor

[zeroexcuses]

Let me see if I understand your argument correctly. Are you saying:

• Rustc does not need to solve algebraic equalities because given the way Rust generics works, these functions are never instantiated with "abstract/symbolic" M/N, but only instantiated with concrete values of M/N, i.e. M = 28, N = 30. And thus, it's trivial for the compiler to see that (M+N)+(M+N) = 116 = (M+M) + (N+N).

If so, this sounds perfectly reasonable (and I have an incorrect model of how Rust generics works).

To test this theory, can you get the following code to compile?

foo<M, N> (m: &Tensor<M>, n: &Tensor<N>) {
  let a: Tensor<M+N> = concat(m, n);
  let b: Tensor<N+M> = concat(n, m);
  let c = a + b; // element wise addition
}

If something like the above compiles, then it's clear that I misunderstood how Rust generics works, with the good news that a lot of 'dimension checking'

• does NOT require solving symbolic equalities over +, -, div by constant, floor, ceil
• and only requires evaluating CONCRETE values over +, -, div by constant, floor, ceil

[jerry73204]

I make a minimal example for your M+N case. https://gist.github.com/jerry73204/86f6bf3848d6cba4b44f32e33a371802

The concrete value is backed by typenum. Look at the source code on docs.rs and it's fairly easy.

[zeroexcuses]

1. You're right. I'm wrong. :-)

2. Thanks for your patience in explaining this step by step.

3. This is really exciting. By the way rustc generics works, it looks like it does not need to solve any "algebraic equalities" and only needs to check "compile time instantiated number equalities."

4. Because of (3), I'm now convinced that: all blas ops should work, convs / pooling should work, rnns/lstms can probably work with batch size specified at compile time. Can you think of any commonly used layer that won't work?

[jerry73204]

Suppose we need the dynamic dimension, such as adjustable batch size. That's where the need of equation solving arises. The project is still in early stage. I would work on dimension sizes known in compile-time, and try to figure out the issue in the mean time.

[LaurentMazare]

Looks pretty cool. The difficulty will indeed probably be in dimension only known at compile time: I would be afraid for type checking to be undecidable in this case as encoding Presburger arithmetic but maybe there can be some tricks around it.

[sherif7810]

Some people are trying to add a similar feature to the Python version.

[fylux]

Has any new progress been made?

[jerry73204]

@fylux It's a busy year for me, and I cannot spare much time in the short term. Nevertheless, I keep it in my radar, so let me summarize the work so far.

Static tensor types is a manner to add more information to compilers. This is the way we achieve dimension checking and inference in compile time. Based on my earlier work in DL, I would desire these features.

• arbitrary sized dimension
• allow dynamic dimension in run-time
• abstraction over storage and back-end

Tensors are not limited to vectors and matrices, and are able to expand. We need to find a way to store arbitrary sized dimension in type. It's possible but it's not a trivial job. By doing so, you need to construct type operators (they are traits itself) to manipulate the dimensions. For example, you may compose two dimension lists during tensor product. AFAIK, the Rust community is still lack of decent work on this kind of advanced types. That why I started off type-freak project. It aims for type-level data structures and control flow (if else, etc).

Dynamic dimension would be a desired feature in our design. Sometimes the batch size is withholded from compiler. In this way, the tensor type would leave some vacant positions that can be only know in runtime. It made our work more complicated. For example , it would be straightforward to add two tensors, if you know their dimensions already. However, what if the dimension of one tensor is only known in run-time? We need to tell the compiler to insert size checking code if needed. Also, the type should be carefully designed to achieve zero-cost property.

I expect the design can be capable with various storage, such as vectors, arrays and something else. I tried to make the type abstract in order that users can extend to their desired back-end. Ultimately, the tensor type only serves the abstraction of additions, products and other ops. We may plug in the TorchStorage to provide actual implementation beneath.

You may visit some of my previous work that are relevant to the idea above.

• type-freak: typed list and type-level control flow and type level algorithms (for example GCD). The master branch is undergoing great refactoring (and halted a bit). You can switch earlier versions to make it compile.
• tch-typed-tensor: The first experimental tensor type design. It proves dimension type can be inferred and checked by compiler. This repo is on hold until the type building blocks become comprehensive.

I also acknowledged works from community. The nalgebra has already done a decent job on matrix type design. It abstracts over various storage and the matrix dimension type is backed on typenum. It's a pity that it know support up to two dimensions. You may also look at frunk. It aims for type level algebraic types.

[jerry73204]

@LaurentMazare @fylux It's been a while and there is some progress so far. Recently I worked on typ that lets you compute types in Rust syntax. Now type-freak, providing type-level data structures and operators, is totally rewritten with typ. It will pave the way to build static typed tensors.

I made an attempt to write dimension type computations with typ. They are in the rust-typed-tensor project.

• matrix multiplication
• convolution shape computation
• dimension concatenation
• PyTorch dimension broadcasting
• tests

I believe the examples show the strong capability to construct static dimensions in Rust. The remaining work is to interoperate with dynamic (in runtime) and static (in compile time) dimensions. That is, whenever the dimension is not known in compile time, we have to dispatch it to runtime computation. I needs next several weeks to figure it out.

Your thoughts and feedback are welcome. Thanks!

[ahirner]

Mind blown @jerry73204. I'm trying to understand the arithmetics behind the rust-typed-tensor tests. Can you explain typ!? I think it generates zero sized Op structs from fns.

If so then:

1. Would such a struct be used to tag Tensor newtypes that carry along some Dims or how else would an integration look like?
2. When and how would other todo! translations come in for typ!?
3. In lieu of expansion, a hand-written minimal example Op would help understanding, I think

Pardon me if these questions are a bit shallow.

[jerry73204]

@ahirner fn inside typ! are translated to traits, not types. The ComputeOp is a type alias to invoke Compute type operator. It can be understood as:

trait Compute<arg1, arg2, ...> { ... } // generated from `fn Compute(...)`
type ComputeOp<arg1, arg2, ...> = <() as Compute<arg1, arg2, ...>>::Output;

It happened that I'm writing a TYP tutorial. You can visit the link to see how it works. https://github.com/jerry73204/typ-book/tree/master/src

As for your questions. Integration can be done in the way below. Just take the dimensions inside input tensors, compute the output dimensions via our type operators, and fill the output type back to tensor.

fn matrix_dot<D1, D2>(lhs: Tensor<D1>, rhs: Tensor<D2>) -> Tensor<MatrixDotOp<D1, D2>>

It takes the advantage of both type safety and type inference. Imagine you have two matrices and call matrix_dot() to compute. The compiler infers the output dimension by expanding MatrixDotOp<D1, D2>, and verifies the return value of matrix_dot() has that shape. The output dimension is inferred by compiler, so it can be regarded as a contract of the tensor API. Also, since the output shape is not given by user, it could reduce some careless errors.

[Narsil]

I attempted my own shot at the issue attempting to use proc_macro instead of pure rust typing here:

https://github.com/Narsil/static_typing_tch

The idea was to actually solve :

foo<M, N> (m: &Tensor<M>, n: &Tensor<N>) {
  let a: Tensor<M+N> = concat(m, n);
  let b: Tensor<N+M> = concat(n, m);
  let c = a + b; // element wise addition
}

And also remove entirely anything dynamic dimensions, as by design you are missing most of the errors there.

By implementing it as a proc_macro we're not bound by the type system at all (but it does imply reimplementing it basically, the good thing is that we only need to reimplement the part we care about. The biggest caveat with the current syntax I chose is that it's not supported to use multiplications or divisions (because they are in traits specifications). It's possible to overcome this to the price of less readable code. The current code I have is this.

#[tensor_check_fn]
fn concat_1(
    left: Tensor<(B, S1), K, D>,
    right: Tensor<(B, S2), K, D>,
) -> Tensor<(B, S1 + S2), K, D> {
    Tensor::cat(&[left, right], 1)
}

This is checked at compile time that Tensor::cat(&[left, right], 1) is actually of shape Tensor<(B, S1 + S2), K, D>

Another example is failing to compile this:

#[tensor_check_fn]
fn add() {
    // The dimensions do not match
    let a = Tensor::ones(&[2, 3], (Kind::Float, Device::Cpu));
    let b = Tensor::ones(&[3, 2], (Kind::Float, Device::Cpu));

    let _c = a + b;  // compilations fails here 
}

The compilation failure shows this:

error: Tensor mismatch Inferred(TensorType { shape: [Value(2), Value(3)], kind: Float, device: Cpu }) is incompatible with Inferred(TensorType { shape: [Value(3), Value(2)], kind: Float, device: Cpu })

This is kind of a weekend project right now, but happy to head feedback on it, maybe find design flaws ahead of time. My current goal is to apply it to a slighly bigger more real piece of code to see if it actually helps.