search

elixir-nx / nx

Multi-dimensional arrays (tensors) and numerical definitions for Elixir
2.55k stars 190 forks source link

Scholar.Neighbors.NNDescent.fit creashes livebook when using EXLA backend #1472

Closed msluszniak closed 2 months ago

msluszniak commented 2 months ago

The following code:

digits_url =
  "https://archive.ics.uci.edu/static/public/887/national+health+and+nutrition+health+survey+2013-2014+(nhanes)+age+prediction+subset.zip"

data =
  Req.get!(digits_url).body

{:ok, [{_, data}]} = :zip.extract(data, [:memory])
df_data = DF.load_csv!(data)

x = df_data |> DF.discard(["age_group", "SEQN"]) |> Nx.stack(axis: 1)
y = Nx.stack(S.cast(df_data["age_group"], :category), axis: 1)

knn = Scholar.Neighbors.NNDescent.fit(x[[0..40, ..]], num_neighbors: 4)

crashes the livebook when backend is set to EXLA

josevalim commented 2 months ago

I was able to isolate it to this:

tensor = Nx.eye({4, 4})
index = Nx.u32(2)
swap = Nx.s64(0)
IO.inspect(tensor[[index, swap]])