Using Nx tensors in type defs

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?


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

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.

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.