Nx - get and set by index for rank1 tensors

iex.exs

iex(1)> t1 = Nx.tensor([3, 4, 5])
#Nx.Tensor<
  s64[3]
  [3, 4, 5]
>
iex(2)> Rank1.at(t1, 0)          
#Nx.Tensor<
  s64
  3
>
iex(3)> t2 = Rank1.put(t1, 0, -2)
#Nx.Tensor<
  s64[3]
  [-2, 4, 5]
>
iex(4)> Rank1.at(t1, 0)          
#Nx.Tensor<
  s64
  3
>
iex(5)> Rank1.at(t2, 0)
#Nx.Tensor<
  s64
  -2
>

rank1.ex

defmodule Rank1 do
  import Nx.Defn

  defn at(tensor, i) do
    shape = Nx.shape(tensor)
    iotas = Nx.iota(shape)
    predicate = Nx.equal(iotas, i)
    values = Nx.select(predicate, tensor, 0)
    Nx.sum(values)
  end

  defn put(tensor, i, val) do
    shape = Nx.shape(tensor)
    iotas = Nx.iota(shape)
    predicate = Nx.equal(iotas, i)
    Nx.select(predicate, val, tensor)
  end
end