search

elixir-nx / nx

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

feat: add shard execution workflow #1557

Open polvalente opened 2 weeks ago

polvalente commented 2 weeks ago

Adds the initial version of the process communication structure for sharded execution.

Does not handle container outputs for the sharded function yet, and also does not yet bring everything together into the compiler jit function.

josevalim commented 1 week ago

Could we fully decouple the workflow definition and execution from Nx? Ideally we would have a workflow like this:

workflow = %{
  0 => %{
    code: &foo(&1, &2, ...),
    args: [1, 2]
  },
  1 =>  %{
    code: &bar(&1),
    args: [2]
  },
  2 => %{
    code: &baz/0,
    args: []
  }
}

And then we pass this to a ProcessExecutor which is completely independent of Nx and tensors. You could also have a Nx executor, but the overall idea is that the Executor should worry about resources and not necessarily tensors (except the resources the tensors are located).