add split function

[tiagodavi]

based on this: elixir-nx/scholar#107

[polvalente]

I believe there's a Nx.Shape.normalize_axis function for doing this

[tiagodavi]

I didn't understand how to use normalize the simple way I would like.

[josevalim]

What happens when size is negative? Do we want to allow negative splitting (which counts from the back)?

[tiagodavi]

not sure if we want this behavior.

[tiagodavi]

@polvalente there's some weird test failing
test/nx/serving_test.exs:637

[polvalente]

@tiagodavi @josevalim I've pushed a refactor to collapse both clauses into the same one, mostly to highlight that the slicing is the same, with the only change being how we calculate it.

I've also changed the examples to operate on a written-out iota instead of arbitrary values because it makes it easier to compare results.

We still need to discuss if we want to accept a negative integer split.

It's easy to do it by setting right after axis_size = : split = if is_integer(split) and split < 0, do: axis_size + split, else: split

[polvalente]

Though I think that negative split might be confusing in the sense that we don't know if the results will be reversed or swapped in the result tuple (either or both)

[josevalim]

@polvalente we will base it on Enum.split (which means looking up from the end).

[josevalim]

But I am fine with postponing this for now. :)

[polvalente]

Right, so given this, my suggestion suffices:

iex(2)> Enum.split(1..10, 10 - 1)
{[1, 2, 3, 4, 5, 6, 7, 8, 9], '\n'}
iex(3)> Enum.split(1..10, -1)
{[1, 2, 3, 4, 5, 6, 7, 8, 9], '\n'

[polvalente]

wdyt?

[tiagodavi]

Awesome, thank you folks.