Suppose I’ve a tensor like this:
f32[3][3][2]
[
[
[0.0, 0.0],
[0.0, 1.0],
[0.0, 2.0]
],
[
[1.0, 0.0],
[1.0, 1.0],
[1.0, 2.0]
],
[
[2.0, 0.0],
[2.0, 1.0],
[2.0, 2.0]
]
]
Is there a way to use the two values of axis 2 as parameters of a function and use the returned value to create a tensor of shape f32[3][3]?
I tried reduce, but it works just iterating along the axis and not just taking the values right?
This will depend a lot on what you’re trying to achieve.
Some functions, such as Nx.sum, Nx.reduce_max and Nx.product accept :axes as options. You can check examples here: https://github.com/elixir-nx/nx/blob/main/nx/guides/tensor-aggregation-101.livemd
However, if you have a more complex function that works in {2}-shaped tensors, and you want to treat your {3, 3, 2} tensor as a collection of those, you could approach it via vectorization: https://github.com/elixir-nx/nx/blob/main/nx/guides/vectorization.livemd
It seems good, but from this point I don’t understand how can I iterate throught my collection of {2}-shaped tensors and call a function for each of them, Nx.map still take alone elements instead the {2}-shaped tensors. I searched in Nx docs but no success.
The idea would for you to implement a function that takes in a {2} shaped tensor, and then vectorize a collection to pass down to it.
For example:
iex(2)> fun = fn x -> Nx.add(x[0], x[1]) end
#Function<42.125776118/1 in :erl_eval.expr/6>
iex(3)> Nx.Defn.jit_apply(fun, [Nx.tensor([10, 20])])
#Nx.Tensor<
s64
30
>
iex(4)> t2 = Nx.tensor([[10, 20], [30, 40], [-10, -20], [-30, -40]])
#Nx.Tensor<
s64[4][2]
[
[10, 20],
[30, 40],
[-10, -20],
[-30, -40]
]
>
iex(5)> Nx.Defn.jit_apply(fun, [t2])
#Nx.Tensor<
s64[2]
[40, 60]
>
iex(6)> Nx.Defn.jit_apply(fun, [Nx.vectorize(t2, :rows)])
#Nx.Tensor<
vectorized[rows: 4]
s64
[30, 70, -30, -70]
>
Notice how vectorizing the input changes how the function reacts to the tensor and yields the sum of each row
Sorry for taking so long, this was a great help, I finally solved my problem. Thank you very much Paulo.
1 Like
