Add support for device kwarg in astype, and add matching utility func #21086

Micky774 · 2024-05-06T16:59:52Z

Towards #20200

This PR adds a private device-placement utility jax._src.numpy.util._place_array to manage array API compliant array placement behavior for use in the jax.numpy namespace. Copies are mediated by the lax._array_copy utility, while device transfer is performed via api.device_put.

cc: @jakevdp

jax/_src/numpy/util.py

tests/lax_numpy_test.py

jax/_src/numpy/util.py

tests/lax_numpy_test.py

jax/_src/numpy/util.py

Micky774 · 2024-05-10T17:35:14Z

@yashk2810 Sorry for the delay in getting this message out. I've updated the PR with simplifications to the _place_array utility, as you suggested. Originally I intended to maximize strict compliance with array API semantics wrt setting copy=False while still specifying device, hence the attempts at detecting when a transfer is not required in the first place (in hindsight, the checks weren't quite right, since I'm still getting familiar with the sharding system).

Is there already a mechanism to predict when two shards are equivalent in the sense that calling device_put(x, shard) would be a no-op if x has sharding equivalent to shard? If so, or if you think something like that is feasible, then it may be worth including here.

cc: @jakevdp if you have any insights regarding how we want to handle copy=False, device=... behavior here

jakevdp · 2024-05-10T17:37:47Z

The easiest thing would just be to error if copy=False and device is specified.

yashk2810 · 2024-05-10T17:59:20Z

Is there already a mechanism to predict when two shards are equivalent in the sense that calling device_put(x, shard) would be a no-op if x has sharding equivalent to shard?

sorry, I don't understand what this means. Also, why do you need such a mechanism?

Micky774 · 2024-05-10T18:07:21Z

sorry, I don't understand what this means. Also, why do you need such a mechanism?

I was looking for a way to tell if x already follows whatever sharding/layout device would specify. That way we could short-cut to a no-op and avoid raising an error in situations like _place_array(x, device=x.sharding, copy=False).

yashk2810 · 2024-05-10T18:10:43Z

I was looking for a way to tell if x already follows whatever sharding/layout device would specify

device_put should do that for you :)

jakevdp · 2024-05-10T19:25:27Z

Yash - the issue is that the semantics of copy=False are "error if a copy is required", so device_put handling that transparently doesn't help because we need to know what's happening in order to know whether to error.

yashk2810 · 2024-05-10T19:28:59Z

Ok, then why not error instead of doing the no-op logic here? If you want to transfer or have it be a no-op should be device_put's job. If device is not None and copy=True, then we should error since it makes no sense right?

jakevdp · 2024-05-10T19:54:21Z

That's what I suggested above: the most conservative thing would be to error if copy and device are used together. But the more "correct" thing would be to somehow introspect whether or not device_put forces a copy in any particular case.

yashk2810 · 2024-05-10T20:04:19Z

We can abstract away a function which can determine that which we can share here but device_put is complex so let's error for now and file a bug against me to give you such a function which you can call here.

jakevdp · 2024-05-24T22:29:25Z

CHANGELOG.md

@@ -37,6 +37,9 @@ Remember to align the itemized text with the first line of an item within a list

 ## jax 0.4.28 (May 9, 2024)

+* New Functionality
+  * {func}`jax.numpy.astype` supports a new `device` keyword argument.
+


Move up to 0.4.29

jakevdp · 2024-05-24T22:33:05Z

jax/experimental/array_api/__init__.py

+Note that JAX may not always strictly adhere to array API device semantics when
+using ``jax.jit``. In particular, specifying the ``device`` argument is
+equivalent to calling ``jax.device_put(x, device)``. For up-to-date details on
+device placement, see the documentation of ``jax.device_put`` for more details.


Usually references go at the bottom of the docstring. You can put this paragraph above the .. _Python... line

yashk2810 · 2024-05-24T23:06:32Z

jax/_src/numpy/util.py

+
+  This utility uses `jax.device_put` for device placement.
+  """
+  out = x


yashk2810 · 2024-05-24T23:08:31Z

jax/_src/numpy/util.py

+      )
+    out = api.device_put(out, device)
+
+  # TODO(micky774): Avoid copy if data has already been copied via device


This todo doesn't make sense? In this branch the device is None. Remove the todo?

yashk2810 · 2024-05-25T04:19:39Z

tests/lax_numpy_test.py

+      device = jax.devices("cpu")[0]
+      expected_sharding = SingleDeviceSharding(device)
+    else:
+      global_mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))


Can we use (2, 2) mesh to coverage across more hardware?

yashk2810 · 2024-05-25T04:20:25Z

tests/lax_numpy_test.py

+      expected_sharding = SingleDeviceSharding(device)
+    else:
+      global_mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
+      device = jax.sharding.NamedSharding(global_mesh, P('x', 'y'))


It would be nice to try other PartitionSpecs too (can you parameterize the test on that too)