Incompatibility between clip and torch.vmap

[TimothyEDawson]

I stumbled into an edge case when trying to apply torch.vmap to some code I had rewritten to utilize array-api-compat. So far everything seems to work just fine, with the exception of clip. Here's a minimal example:

import array_api_compat.torch as xp
import torch


def apply_clip(a):
    return torch.clip(a, min=0, max=30)


def apply_clip_compat(a):
    return xp.clip(a, min=0, max=30)


a = xp.asarray([[5.1, 2.0, 64.1, -1.5]])

print(apply_clip(a))
print(apply_clip_compat(a))

v1 = torch.vmap(apply_clip)
print(v1(a))

v2 = xp.vmap(apply_clip_compat)
print(v2(a))

Which raises the following error:

[user@domain ~]$ python test_clip.py 
tensor([[ 5.1000,  2.0000, 30.0000,  0.0000]])
tensor([[ 5.1000,  2.0000, 30.0000,  0.0000]])
tensor([[ 5.1000,  2.0000, 30.0000,  0.0000]])
Traceback (most recent call last):
  File "test_clip.py", line 22, in <module>
    print(v2(a))
          ~~^^^
  File ".venv/lib/python3.13/site-packages/torch/_functorch/apis.py", line 202, in wrapped
    return vmap_impl(
        func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs
    )
  File ".venv/lib/python3.13/site-packages/torch/_functorch/vmap.py", line 334, in vmap_impl
    return _flat_vmap(
        func,
    ...<6 lines>...
        **kwargs,
    )
  File ".venv/lib/python3.13/site-packages/torch/_functorch/vmap.py", line 484, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "test_clip.py", line 10, in apply_clip_compat
    return xp.clip(a, min=0, max=30)
           ~~~~~~~^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.13/site-packages/array_api_compat/_internal.py", line 35, in wrapped_f
    return f(*args, xp=xp, **kwargs)
  File ".venv/lib/python3.13/site-packages/array_api_compat/common/_aliases.py", line 424, in clip
    out[()] = x
    ~~~^^^^
RuntimeError: vmap: inplace arithmetic(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `self` not being vmapped over in a vmap. Please try to use out-of-place operators instead of inplace arithmetic. If said operator is being called inside the PyTorch framework, please file a bug report instead.

I totally understand if full support for torch.vmap is out of scope, but figured it might be worth raising the issue in case there's something which requires fixing.

[lucascolley]

that is a shame!

array-api-compat/array_api_compat/common/_aliases.py

Lines 438 to 439 in be9eff7

	# Return a scalar for 0-D
	return out[()]

For most things, [()] works to get NumPy scalars for NumPy, and as a no-op for other arrays. But with vmap it seems the no-op assumption breaks down.

In the long-run, I think this calls for an as_scalar_if_numpy helper to deal with the pesky NumPy special-case in a safe way. Not sure if there is a simpler workaround in the short term. WDYT @ev-br ?

[ev-br]

In array-api-compat specifically, I think it's a matter of a simple refactoring to move the [()] out of common/ and into "numpy/" and other specific folders.

EDIT: one other offender is common/sign.

[lucascolley]

Good shout!

[ev-br]

It would be great to understand better the origin of this error.
If I replace out[()] = x with out = torch.clone(x) at

array-api-compat/array_api_compat/common/_aliases.py

Line 424 in be9eff7

out[()] = x

, it immediately fails a few lines below at

array-api-compat/array_api_compat/common/_aliases.py

Line 430 in be9eff7

out[ia] = a[ia]

While this all can be worked around, it begs the question what is the deeper reason behind all this. Surely it cannot be that indexing operations are incompatible with vmap?

[ev-br]

Ah, OK, This seems to be the answer: https://discuss.pytorch.org/t/vmap-inplace-arithmetic-error/196556/2

you can’t have a write to an existing tensor

So indeed mutating indexing operations are incompatible with torch.vmap. Which means array_api_compat.torch.clip should be rewritten in terms of torch.clamp instead of emulating it as it does ATM.

[ev-br]

A tentative fix is in #353, but I'm not sure it's entirely correct because of promotion rules. Thoughts?