put(): expand over repeat internally

Also test 0d indices

Author: honno

torch_np/_funcs_impl.py

@@ -901,14 +901,12 @@ def put(
     mode: NotImplementedType = "raise",
 ):
     v = v.type(a.dtype)
-    # If ind is larger than v, broadcast v to the would-be resulting shape. Any
+    # If ind is larger than v, expand v to at least the size of ind. Any
     # unnecessary trailing elements are then trimmed.
     if ind.numel() > v.numel():
         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
+        v = v.unsqueeze(0).expand((ratio,) + v.shape)
+    # Trim unnecessary elements, regarldess if v was expanded or not. Note
     # np.put() trims v to match ind by default too.
     if ind.numel() < v.numel():
         v = v.flatten()

torch_np/tests/test_xps.py

@@ -117,9 +117,12 @@ def test_put(np_x, data):
     tnp_x = tnp.asarray(np_x.copy()).astype(np_x.dtype.name)
 
     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)
-    )
+    if result_shape == ():
+        ind_strat = st.integers(np_x.size)
+    else:
+        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(