Combine backend and device arguments
#1898
Conversation
|
This looks promising! What's the status? Want to get it reviewed, or is there more work to be done? |
| @@ -107,8 +107,8 @@ def jit(fun, static_argnums=(), device=None, backend=None): | |||
| Optional, the Device the jitted function will run on. (Available devices | |||
| can be retrieved via ``jax.devices()``.) The default is inherited from | |||
| XLA's DeviceAssignment logic and is usually to use ``jax.devices()[0]``. | |||
| backend: This is an experimental feature and the API is likely to change. | |||
| Optional, a string representing the xla backend. 'cpu','gpu', or 'tpu'. | |||
| Can also be passed ``'cpu'``, ``'gpu'``, or ``'tpu'`` to use the default | |||
There was a problem hiding this comment.
It does not seem nice to overload the device argument with a string denoting the platform. I assume this would be rarely needed, why not ask the user to use devices=devices("gpu")[0]?
| @@ -107,8 +107,8 @@ def jit(fun, static_argnums=(), device=None, backend=None): | |||
| Optional, the Device the jitted function will run on. (Available devices | |||
| can be retrieved via ``jax.devices()``.) The default is inherited from | |||
| XLA's DeviceAssignment logic and is usually to use ``jax.devices()[0]``. | |||
| backend: This is an experimental feature and the API is likely to change. | |||
There was a problem hiding this comment.
I think that we should get in the habit of deprecating arguments from public APIs before removing them. I would suggest keeping the backend, immediately turning it into device=devices(backend)[0], issuing in awarning, adding the note to Changelog, adding a reminder somewhere to remove it in a future release. This is a lot more expensive than just removing it, but I think that as we have more and more users it is a good thing to do.
Also, it may make sense to do codesearch/ for people who use the backend argument
| backend: This is an experimental feature and the API is likely to change. | ||
| Optional, a string representing the xla backend. 'cpu','gpu', or 'tpu'. | ||
| device: This is an experimental feature and the API is likely to change. | ||
| Optional, either a Device object or ``'cpu'``, ``'gpu'``, or ``'tpu'``. |
There was a problem hiding this comment.
If we do remove the string device arguments, we should fix this docstring too
|
|
||
| def canonicalize_device_arg(device_arg): | ||
| if isinstance(device_arg, str): | ||
| if device_arg not in VALID_PLATFORMS: |
There was a problem hiding this comment.
Maybe check instead whether it is a supported platform.
| @@ -48,7 +48,7 @@ class MultiBackendTest(jtu.JaxTestCase): | |||
| def testMultiBackend(self, backend): | |||
| if backend not in ('cpu', jtu.device_under_test(), None): | |||
| raise SkipTest() | |||
| @partial(api.jit, backend=backend) | |||
There was a problem hiding this comment.
Instead of doing this, I would rewrite (or eliminate) the test. It is not clear that in the absence of the backend argument this test adds value.
|
Closing this PR. Whatever its merits were, it is quite stale. |
No description provided.