LaurentMazare / tch-rs

Rust bindings for the C++ api of PyTorch.
Apache License 2.0
4.39k stars 347 forks source link

Shapes expect i64 instead of usize #51

Open grtlr opened 5 years ago

grtlr commented 5 years ago

Maybe this comes from generating the bindings automatically, but is there a use case for expecting &[i64] in Tensor::reshape and similar methods? Maybe it would make sense to change this to &[usize]? This change would also bring Tensor more in line with ndarray, which uses usize to represent shapes.

LaurentMazare commented 5 years ago

Right, having usize at some places would be better. That said, a commonly used value is -1 both for indexes and shapes. For indexes it means start from the end, -1 would be the last value, -2 the value before and so on. For shapes it's used for a single dimension and gets an appropriate value based on the tensor total number of elements (this number of element has to be divisible by the product of sizes for other dimensions of course). E.g. one could write a.reshape([batch_size, -1]) to flatten a. This explains why i64 are used there, it's less verbose than using Option<usize> for shapes or some other alternative for indexes, it's also close to how the python api works but it certainly doesn't feel very rust idiomatic.

grtlr commented 5 years ago

Thanks for the quick reply! I like the idea of tch-rs staying close to the Python API. For now I will just start to write some more idiomatic Rust wrappers around some of these methods for my own needs. Maybe in the future we can have a separate crate which handles these kind of things. Such a crate was briefly discussed in #12, do you have any ideas in mind here?

LaurentMazare commented 5 years ago

I don't have much concrete ideas but maybe a couple of these would be good:

I would have hoped for the first two points to be doable when generating the binding code but I haven't found enough information in the Declarations.yaml file that is used for generation. Fwiw tensorflow provides some description of operations with all the details on the supported types which make it possible to generate better type annotations, e.g. for conv2d here is some generated ocaml code: conv2d can only apply to tensors of float or double and return a tensor of the same type.

Also if you get some interesting wrappers that you think should be part of tch-rs let me know via a PR or an issue.

LaurentMazare commented 5 years ago

Thinking more about this, it would certainly be nice to use more of usize for dimensions/shape, e.g. for size or for the various nn creators. The only exception I can think of are view, get, narrow which should handle negative indexing and values. I'll try making some changes for this when I get a chance.