Repository: numba/numba

Issue:
Cannot resolve setitem for array slice
The documentation leads me to believe this should work. Section 2.6.2.1: "Full basic indexing and slicing is supported." This is Python 3.5, NumPy 1.10.4, and Numba 0.25.

``` python
import numpy as np
import numba
def test_func():
    a = np.ones(2, dtype=np.int32)
    b = np.zeros((2, 2), dtype=np.int32)
    b[:] = a
    return b
numba.jit(nopython=True)(test_func)()
```

```
TypingError: Failed at nopython (nopython frontend)
Cannot resolve setitem: array(int32, 2d, C)[slice(None, None, None)] = array(int32, 1d, C)
```

Workaround:

``` python
import numpy as np
import numba
def test_func():
    a = np.ones(2, dtype=np.int32)
    b = np.zeros((2, 2), dtype=np.int32)
    for idx in range(b.shape[0]):
        b[idx, :] = a
    return b
numba.jit(nopython=True)(test_func)()
```

```
array([[1, 1],
       [1, 1]], dtype=int32)
```



Discussion:
This looks like we failed to implement the broadcast rules for slicing completely.  We'll take a look.


Resolving patch:
diff --git a/numba/cgutils.py b/numba/cgutils.py
--- a/numba/cgutils.py
+++ b/numba/cgutils.py
@@ -71,6 +71,21 @@ def create_struct_proxy(fe_type, kind='value'):
     return res
 
 
+def copy_struct(dst, src, repl={}):
+    """
+    Copy structure from *src* to *dst* with replacement from *repl*.
+    """
+    repl = repl.copy()
+    # copy data from src or use those in repl
+    for k in src._datamodel._fields:
+        v = repl.pop(k, getattr(src, k))
+        setattr(dst, k, v)
+    # use remaining key-values in repl
+    for k, v in repl.items():
+        setattr(dst, k, v)
+    return dst
+
+
 class _StructProxy(object):
     """
     Creates a `Structure` like interface that is constructed with information
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -5,7 +5,6 @@
 
 from __future__ import print_function, absolute_import, division
 
-import collections
 import functools
 import math
 
@@ -1181,6 +1180,88 @@ def src_cleanup():
     return src_getitem, src_cleanup
 
 
+def _bc_adjust_dimension(context, builder, shapes, strides, target_shape):
+    """
+    Preprocess dimension for broadcasting.
+    Returns (shapes, strides) such that the ndim match *target_shape*.
+    When expanding to higher ndim, the returning shapes and strides are
+    prepended with ones and zeros, respectively.
+    When truncating to lower ndim, the shapes are checked (in runtime).
+    All extra dimension must have size of 1.
+    """
+    zero = context.get_constant(types.uintp, 0)
+    one = context.get_constant(types.uintp, 1)
+
+    # Adjust for broadcasting to higher dimension
+    if len(target_shape) > len(shapes):
+        nd_diff = len(target_shape) - len(shapes)
+        # Fill missing shapes with one, strides with zeros
+        shapes = [one] * nd_diff + shapes
+        strides = [zero] * nd_diff + strides
+    # Adjust for broadcasting to lower dimension
+    elif len(target_shape) < len(shapes):
+        # Accepted if all extra dims has shape 1
+        nd_diff = len(shapes) - len(target_shape)
+        dim_is_one = [builder.icmp_unsigned('==', sh, one)
+                      for sh in shapes[:nd_diff]]
+        accepted = functools.reduce(builder.and_, dim_is_one,
+                                    cgutils.true_bit)
+        # Check error
+        with builder.if_then(builder.not_(accepted), likely=False):
+            msg = "cannot broadcast source array for assignment"
+            context.call_conv.return_user_exc(builder, ValueError, (msg,))
+        # Truncate extra shapes, strides
+        shapes = shapes[nd_diff:]
+        strides = strides[nd_diff:]
+
+    return shapes, strides
+
+
+def _bc_adjust_shape_strides(context, builder, shapes, strides, target_shape):
+    """
+    Broadcast shapes and strides to target_shape given that their ndim already
+    matches.  For each location where the shape is 1 and does not match the
+    dim for target, it is set to the value at the target and the stride is
+    set to zero.
+    """
+    bc_shapes = []
+    bc_strides = []
+    zero = context.get_constant(types.uintp, 0)
+    one = context.get_constant(types.uintp, 1)
+    # Adjust all mismatching ones in shape
+    mismatch = [builder.icmp_signed('!=', tar, old)
+                for tar, old in zip(target_shape, shapes)]
+    src_is_one = [builder.icmp_signed('==', old, one) for old in shapes]
+    preds = [builder.and_(x, y) for x, y in zip(mismatch, src_is_one)]
+    bc_shapes = [builder.select(p, tar, old)
+                 for p, tar, old in zip(preds, target_shape, shapes)]
+    bc_strides = [builder.select(p, zero, old)
+                  for p, old in zip(preds, strides)]
+    return bc_shapes, bc_strides
+
+
+def _broadcast_to_shape(context, builder, arrtype, arr, target_shape):
+    """
+    Broadcast the given array to the target_shape.
+    Returns (array_type, array)
+    """
+    # Compute broadcasted shape and strides
+    shapes = cgutils.unpack_tuple(builder, arr.shape)
+    strides = cgutils.unpack_tuple(builder, arr.strides)
+
+    shapes, strides = _bc_adjust_dimension(context, builder, shapes, strides,
+                                           target_shape)
+    shapes, strides = _bc_adjust_shape_strides(context, builder, shapes,
+                                               strides, target_shape)
+    new_arrtype = arrtype.copy(ndim=len(target_shape), layout='A')
+    # Create new view
+    new_arr = make_array(new_arrtype)(context, builder)
+    repl = dict(shape=cgutils.pack_array(builder, shapes),
+                strides=cgutils.pack_array(builder, strides))
+    cgutils.copy_struct(new_arr, arr, repl)
+    return new_arrtype, new_arr
+
+
 def fancy_setslice(context, builder, sig, args, index_types, indices):
     """
     Implement slice assignment for arrays.  This implementation works for
@@ -1202,13 +1283,15 @@ def fancy_setslice(context, builder, sig, args, index_types, indices):
     if isinstance(srcty, types.Buffer):
         # Source is an array
         src_dtype = srcty.dtype
+        index_shape = indexer.get_shape()
         src = make_array(srcty)(context, builder, src)
+        # Broadcast source array to shape
+        srcty, src = _broadcast_to_shape(context, builder, srcty, src, index_shape)
         src_shapes = cgutils.unpack_tuple(builder, src.shape)
         src_strides = cgutils.unpack_tuple(builder, src.strides)
         src_data = src.data
 
         # Check shapes are equal
-        index_shape = indexer.get_shape()
         shape_error = cgutils.false_bit
         assert len(index_shape) == len(src_shapes)
 
diff --git a/numba/typing/arraydecl.py b/numba/typing/arraydecl.py
--- a/numba/typing/arraydecl.py
+++ b/numba/typing/arraydecl.py
@@ -183,13 +183,11 @@ def generic(self, args, kws):
         if isinstance(res, types.Array):
             # Indexing produces an array
             if isinstance(val, types.Array):
-                if (val.ndim == res.ndim and
-                    self.context.can_convert(val.dtype, res.dtype)):
-                    # Allow assignement of same-dimensionality compatible-dtype array
-                    res = val
-                else:
-                    # NOTE: array broadcasting is unsupported
+                if not self.context.can_convert(val.dtype, res.dtype):
+                    # DType conversion not possible
                     return
+                else:
+                    res = val
             elif isinstance(val, types.Sequence):
                 if (res.ndim == 1 and
                     self.context.can_convert(val.dtype, res.dtype)):


Test patch:
diff --git a/numba/tests/test_indexing.py b/numba/tests/test_indexing.py
--- a/numba/tests/test_indexing.py
+++ b/numba/tests/test_indexing.py
@@ -131,6 +131,12 @@ def empty_tuple_usecase(a):
 def setitem_usecase(a, index, value):
     a[index] = value
 
+
+@njit
+def setitem_broadcast_usecase(a, value):
+    a[:] = value
+
+
 def slicing_1d_usecase_set(a, b, start, stop, step):
     a[start:stop:step] = b
     return a
@@ -967,11 +973,86 @@ def test_setitem(self):
         # Using a 0-d array as scalar index
         setitem_usecase(arr, np.array(3).astype(np.uint16), 8)
         self.assertEqual(arr.tolist(), [0, 42, 2, 8, 4])
-        # Broadcasting
+        # Scalar Broadcasting
         arr = np.arange(9).reshape(3, 3)
         setitem_usecase(arr, 1, 42)
         self.assertEqual(arr.tolist(), [[0, 1, 2], [42, 42, 42], [6, 7, 8]])
 
+    def test_setitem_broadcast(self):
+        """
+        broadcasted array assignment
+        """
+        # Scalar Broadcasting
+        dst = np.arange(5)
+        setitem_broadcast_usecase(dst, 42)
+        self.assertEqual(dst.tolist(), [42] * 5)
+        # 1D -> 2D Array Broadcasting
+        dst = np.arange(6).reshape(2, 3)
+        setitem_broadcast_usecase(dst, np.arange(1, 4))
+        self.assertEqual(dst.tolist(), [[1, 2, 3], [1, 2, 3]])
+        # 2D -> 2D Array Broadcasting
+        dst = np.arange(6).reshape(2, 3)
+        setitem_broadcast_usecase(dst, np.arange(1, 4).reshape(1, 3))
+        self.assertEqual(dst.tolist(), [[1, 2, 3], [1, 2, 3]])
+        # 2D -> 4D Array Broadcasting
+        dst = np.arange(12).reshape(2, 1, 2, 3)
+        setitem_broadcast_usecase(dst, np.arange(1, 4).reshape(1, 3))
+        inner2 = [[1, 2, 3], [1, 2, 3]]
+        self.assertEqual(dst.tolist(), [[inner2]] * 2)
+        # 2D -> 1D Array Broadcasting
+        dst = np.arange(5)
+        setitem_broadcast_usecase(dst, np.arange(1, 6).reshape(1, 5))
+        self.assertEqual(dst.tolist(), [1, 2, 3, 4, 5])
+        # 4D -> 2D Array Broadcasting
+        dst = np.arange(6).reshape(2, 3)
+        setitem_broadcast_usecase(dst, np.arange(1, 1 + dst.size).reshape(1, 1, 2, 3))
+        self.assertEqual(dst.tolist(), [[1, 2, 3], [4, 5, 6]])
+
+    def test_setitem_broadcast_error(self):
+        # higher dim assigned into lower dim
+        # 2D -> 1D
+        dst = np.arange(5)
+        src = np.arange(10).reshape(2, 5)
+        with self.assertRaises(ValueError) as raises:
+            setitem_broadcast_usecase(dst, src)
+        errmsg = str(raises.exception)
+        self.assertEqual('cannot broadcast source array for assignment',
+                         errmsg)
+        # 3D -> 2D
+        dst = np.arange(5).reshape(1, 5)
+        src = np.arange(10).reshape(1, 2, 5)
+        with self.assertRaises(ValueError) as raises:
+            setitem_broadcast_usecase(dst, src)
+        errmsg = str(raises.exception)
+        self.assertEqual('cannot assign slice from input of different size',
+                         errmsg)
+        # lower to higher
+        # 1D -> 2D
+        dst = np.arange(10).reshape(2, 5)
+        src = np.arange(4)
+        with self.assertRaises(ValueError) as raises:
+            setitem_broadcast_usecase(dst, src)
+        errmsg = str(raises.exception)
+        self.assertEqual('cannot assign slice from input of different size',
+                         errmsg)
+
+    def test_slicing_1d_broadcast(self):
+        # 1D -> 2D sliced (1)
+        dst = np.arange(6).reshape(3, 2)
+        src = np.arange(1, 3)
+        slicing_1d_usecase_set(dst, src, 0, 2, 1)
+        self.assertEqual(dst.tolist(), [[1, 2], [1, 2], [4, 5]])
+        # 1D -> 2D sliced (2)
+        dst = np.arange(6).reshape(3, 2)
+        src = np.arange(1, 3)
+        slicing_1d_usecase_set(dst, src, 0, None, 2)
+        self.assertEqual(dst.tolist(), [[1, 2], [2, 3], [1, 2]])
+        # 2D -> 2D sliced (3)
+        dst = np.arange(6).reshape(3, 2)
+        src = np.arange(1, 5).reshape(2, 2)
+        slicing_1d_usecase_set(dst, src, None, 2, 1)
+        self.assertEqual(dst.tolist(), [[1, 2], [3, 4], [4, 5]])
+
     def test_setitem_readonly(self):
         arr = np.arange(5)
         arr.flags.writeable = False