ENH: add new function one_hot

[NeilGirdhar]

Fixes #305

Questions:

• What should the default output dtype be? Currently, it's the default float type on Jax/PyTorch, or else float64. Should I try to get the default float type on the platform: xp.empty(()).dtype?

[lucascolley]

RE dtype, I think data-apis/array-api#848 will give us something slightly cleaner down the line.

In SciPy we have been using https://github.com/scipy/scipy/blob/main/scipy/_lib/_array_api.py#L399 with force_floating=True. Taking xp.empty(()).dtype is probably fine for now though.

[NeilGirdhar]

Yeah, I'm not 100% sure how those links will help in this case? We're not casting one thing to another kind, or promoting to float. We just want the default float dtype irrespective of what was passed in. You may want to consider giving names to:

xp.asarray(1j).dtype  # Default complex
xp.asarray(1).dtype  # Default int
xp.empty(()).dtype  # Default float

I'm not sure though, and these are one-liners.

[lucascolley]

I'm not 100% sure how those links will help in this case?

the point is that instead of writing

dtype = xp.empty(()).dtype
x = xp.zeros(..., dtype=dtype)

we would have

x = xp.zeros(...)
x = xp.astype(x, 'real floating')

[NeilGirdhar]

Ah, right. Okay. I guess you mean askind, right?

Also, could you help me solve the Dask errors? This is all foreign to me.

And how do I make an array with a non-concrete size? (x.size=None)

[lucascolley]

I guess you mean askind, right?

Nope, data-apis/array-api#848 suggests overloading the dtype parameter of astype for this. Feel free to comment over there if you disagree!

[NeilGirdhar]

suggests overloading the dtype parameter of astype for this. Feel free to comment over there if you disagree!

Oh, no worries, I don't have a strong opinion. Just trying to keep up with all the planned changes 😄

Did you see my edits? I could use some guidance with the Dask errors.

[crusaderky]

• What should the default output dtype be? Currently, it's the default float type on Jax/PyTorch, or else float64. Should I try to get the default float type on the platform: xp.empty(()).dtype?

to me this doesn't make much sense. Why shouldn't it be bool?

[NeilGirdhar]

• What should the default output dtype be? Currently, it's the default float type on Jax/PyTorch, or else float64. Should I try to get the default float type on the platform: xp.empty(()).dtype?

to me this doesn't make much sense. Why shouldn't it be bool?

I'd rather follow what the libraries are doing than force double conversion for delegated code. If you do that, most people would end up having to write their own one_hot method in order to avoid it.

In general, the reason it's not bool is because these values often serve as the inputs to machine learning algorithms.

[NeilGirdhar]

@lucascolley Ready for your review

[crusaderky]

[EDIT] just realised that I have had an unsent review hanging for the last 3 days. Apologies.

What about x[..., None] == xp.arange(num_classes, dtype=x.dtype, device=_compact.device(x))?

that should be all you need to do to implement one_hot in a way that is both performant and Array API compliant.
(the numpy special case may be faster; needs benchmarking).

[crusaderky]

I think with the new algorithm none sizes are supported (needs a test though)

[crusaderky]

Suggested change

assert is_jax_array(x)

We typically don't do this. It's needlessly expensive.

[NeilGirdhar]

I agree that we should be consistent with the rest of the package, so I'll remove it. However, you may want to revisit this policy since with jitted code such calls are basically free. This is true for Jax and Pytorch, which are the backends that I added assertions with. Also, assertions disappear when optimizations are turned on.

The assertions have the benefit of improve type checking.

[crusaderky]

This is almost ready to go; a few nits below.

It's worth noting that the current implementation makes it impossible to write

symbols, idx = xp.unique_inverse(x)
xpx.one_hot(idx, symbols.size)

as a pattern to build a one-hot map of arbitrary symbols on Dask, unless you know in advance the maximum number of unique symbols.
This could be fixed in a follow-up but I'm unsure about real-life interest in it.

[lucascolley]

looks like there was a rebase hiccup

[lucascolley]

thanks both, looks great!

[NeilGirdhar]

Thanks @lucascolley and @crusaderky for the thorough and quick reviews!

[crusaderky]

Only one nit; see above

[lucascolley]

looks like there was a rebase hiccup

[lucascolley]

thanks again!