[BUG] in-place updating of array slice unexpectedly fails due to broadcasting problem

[astrotuna201]

Describe the bug mlx fails to update a slice of an array, equivalent in numpy works (surfaced in mlx-swift, checked with python/numpy as per below) suspected bug in mlx::core

Edit: perhaps related to #1015

To Reproduce

• in mlx:

  import mlx.core as mx
  a = mx.ones([2,6,6,6])
  b = mx.zeros([3,4,4,4])
  b[0,0:4,3,0:4] = a[0,1:5, 5,1:5]
  # ValueError: Cannot broadcast array of shape (4,4) into shape (1,4,1,4).

Expected behavior

• in numpy:

  import numpy as np
  a = np.ones([2,6,6,6])
  b = np.zeros([3,4,4,4])
  b[0,0:4,3,0:4] = a[0,1:5, 5,1:5] #works
  np.count_nonzero(b == 1) #16
  np.count_nonzero(b == 0) #176

Desktop (please complete the following information):

• OS Version: MacOS 14.4.1
  • Version: Installing collected packages: numpy, mlx Successfully installed mlx-0.12.0 numpy-1.26.4

Additional context This slice update also fails in mlx-swift.

import MLX
let a = MLXArray.ones([2,6,6,6])
let b = MLXArray.zeros([3,4,4,4])
b[0,0..<4, 3, 0..<4] = a[0, 1..<5, 5, 1..<5]

Backtrace:

std::invalid_argument: Cannot broadcast array of shape (4,4) into shape (1,4,1,4).

  * frame #0: 0x0000000108b1b7a6 MLXTests`mlx::core::broadcast_to(a=0x00006000015abeb0, shape=size=4, s= Active Type = mlx::core::Stream ) at ops.cpp:1105:5
    frame #1: 0x0000000108b206f0 MLXTests`mlx::core::slice_update(src=0x00006000015abbb0, update=0x00006000015abeb0, start=size=4, stop=size=4, strides=size=4, s= Active Type = mlx::core::Stream ) at ops.cpp:588:29
    frame #2: 0x0000000107fdb1ed MLXTests`mlx_slice_update(src=0x00006000015abba0, update=0x00006000015abea0, start=0x0000600001bc40e0, num_start=4, stop=0x0000600001bc4140, num_stop=4, strides=0x0000600001bc4170, num_strides=4, s=0x00006000015abe80) at ops.cpp:750:10
    frame #3: 0x0000000107f13bc8 MLXTests`updateSlice(src=(ctx = 0x00006000015abba0), operations=4 values, update=(ctx = 0x00006000015abea0), stream=MLX.StreamOrDevice @ 0x00007ff7bfefc420) at MLXArray+Indexing.swift:911:9

in frame #3:
po src.shape
▿ 4 elements
  - 0 : 3
  - 1 : 4
  - 2 : 4
  - 3 : 4
(lldb) po update.shape
▿ 2 elements
  - 0 : 4
  - 1 : 4
(lldb) po starts
▿ 4 elements
  - 0 : 0
  - 1 : 0
  - 2 : 3
  - 3 : 0
(lldb) po ends
▿ 4 elements
  - 0 : 1
  - 1 : 4
  - 2 : 4
  - 3 : 4
(lldb) po strides
▿ 4 elements
  - 0 : 1
  - 1 : 1
  - 2 : 1
  - 3 : 1

[awni]

This is fixed in #1035, we should file a companion issue in MLX Swift.

[awni]

Closing this as dup and opened a issue for MLX Swift: https://github.com/ml-explore/mlx-swift/issues/76