put(): support and test broadcast-incompatible args

Author: honno

torch_np/_funcs_impl.py

@@ -904,8 +904,10 @@ def put(
     # If ind is larger than v, broadcast v to the would-be resulting shape. Any
     # unnecessary trailing elements are then trimmed.
     if ind.numel() > v.numel():
-        result_shape = torch.broadcast_shapes(v.shape, ind.shape)
-        v = torch.broadcast_to(v, result_shape)
+        ratio = (ind.numel() + v.numel() - 1) // v.numel()
+        sizes = [ratio]
+        sizes.extend([1 for _ in range(v.dim() - 1)])
+        v = v.repeat(*sizes)
     # Trim unnecessary elements, regarldess if v was broadcasted or not. Note
     # np.put() trims v to match ind by default too.
     if ind.numel() < v.numel():

torch_np/tests/test_xps.py

@@ -4,6 +4,7 @@
 These tests aren't specifically for testing Array API adoption!
 """
 import cmath
+import math
 import warnings
 
 import pytest
@@ -115,11 +116,20 @@ def test_put(np_x, data):
 
     tnp_x = tnp.asarray(np_x.copy()).astype(np_x.dtype.name)
 
-    result_shapes = st.shared(nps.array_shapes())
-    ind = data.draw(
-        nps.integer_array_indices(np_x.shape, result_shape=result_shapes), label="ind"
+    result_shape = data.draw(nps.array_shapes(), label="result_shape")
+    ind_strat = nps.integer_array_indices(
+        np_x.shape, result_shape=st.just(result_shape)
+    )
+    ind = data.draw(ind_strat | ind_strat.map(np.asarray), label="ind")
+    v = data.draw(
+        nps.arrays(
+            dtype=np_x.dtype,
+            shape=nps.array_shapes().filter(
+                lambda s: math.prod(s) > math.prod(result_shape)
+            ),
+        ),
+        label="v",
     )
-    v = data.draw(nps.arrays(dtype=np_x.dtype, shape=result_shapes), label="v")
 
     tnp_x_copy = tnp_x.copy()
     np.put(np_x, ind, v)