Using Nx tensors in type defs

Elixir Framework Forums Nx Forum
, ,
1

I’d like to use typespecs to document that certain functions take one of a few different tensors with specific shape/values like

@type my_tensors() :: Nx.tensor([0, 1]) | Nx.tensor([1, 0]) | Nx.tensor([-1, 0]) | Nx.tensor([0, -1])

The compiler doesn’t like this giving me an error about “unexpected list in typespec”. Is there any way around this?

2 Likes
2

I don’t see anything about literal lists in the typespec docs at all - the only supported literals are atoms / numbers / binaries.

Beyond that, there’s no Nx.tensor type in Nx either…

1 Like
3

Nx.Tensor defines the struct that is returned by Nx.tensor. I think I want @type to be a way to create an equivalent to rust’s Enums or custom types in elm, but it’s just not that kind of tool.

4

It can do some of the same things, with the right kind of literals. For instance, you could define types like in that example:

@type user_status :: :regular | :visitor
@type user :: %{status: user_status, name: binary}

@type user_with_extra :: {:regular, binary, integer} | {:visitor, binary}

The main restriction is that while you can write types that take arguments, they can’t do arithmetic or pattern-match and live in a separate namespace from plain functions.

On that note, there’s no Nx.tensor @type - the closest thing is Nx.Tensor.t, but that uses the argument to specify the backend type not the shape.