ENH: implement at
#53
Conversation
crusaderky
force-pushed
the
at
branch
9 times, most recently
from
376ad49 to
4a0943a
Compare
|
Rebased on top of #58 |
|
|
Any idea how to fix this? Looks like the ubuntu_latest VM has an obsolete driver (or more likely no driver)
|
|
I don't think we can get GPU CI without paying someone for it, cc @rgommers . |
Can cupy run on a CPU-only host? |
Yep, that's on my radar to push forward this month, on multiple projects. Please feel free to open a new issue and assign it to me. I think we can hook up a shared GPU runner between this project and
I don't think so. |
|
I gave a final round of polish. This is ready for final review and approval. |
|
|
I resolved the merge conflicts after adding |
|
@lucascolley if you don't mind building against array-api-compat git tip until their next release, this PR is ready for final review and merge |
There was a problem hiding this comment.
@lucascolley if you don't mind building against array-api-compat git tip until their next release, this PR is ready for final review and merge
That's fine with me, but might be cleaner to get an array-api-compat release out first unless there are any big blockers.
Let's give @rgommers and @jakevdp time to take a look if they would like to.
Once (or before) this is merged, do you think you could make a PR to my branch at scikit-learn/scikit-learn#30340? I think we will have to transition scikit-learn from using array-api-compat as an optional dependency to vendoring it, but that sounds feasible based on scikit-learn/scikit-learn#30367 (comment).
|
| if res is not None: | ||
| return res | ||
| assert x is not None | ||
| x[self._idx] = elwise_op(x[self._idx], y) |
There was a problem hiding this comment.
As currently implemented, the _iop path has different semantics for repeated indices between JAX and NumPy:
>>> import numpy as np
>>> x = np.zeros(4)
>>> idx = np.array([1, 2, 2, 3, 3, 3])
>>> x[idx] = x[idx] + 1
>>> x
array([0., 1., 1., 1.])
>>> import jax.numpy as jnp
>>> x = jnp.zeros(4)
>>> idx = jnp.array([1, 2, 2, 3, 3, 3])
>>> x.at[idx].add(1)
Array([0., 1., 2., 3.], dtype=float32)At the very least, the difference should be documented.
There was a problem hiding this comment.
it's already documented at lines 612:622
Implement a new
at(x, idx)orat(x)[idx]function, mocking the syntax of JAX's omonymous method .This is propaedeutic to JAX support in libraries that support the Array API, e.g. scipy.
Moved from data-apis/array-api-compat#205
Blockers