Choosing values from an mlx array based on indices and booleans with different input and output sizes

[kyrollosyanny]

Hello,

I am trying to write a custom image resizing algorithm in MLX. For example, let's try to do bilinear interpolation. One problem I am running into is that MLX cannot use booleans to access array information. For example, after I figure out how the output pixels map to the input pixels, we need to assign the output pixel to input pixel based on some indices. I think there might be a combination of masks and mx.where to accomplish that but I am not sure how to go about it?

Here is an example with a simple case, last line gives an error.

original_img=mx.array(I)  ## input image to resize
old_h, old_w, c = original_img.shape
###new image size
new_h=300
new_w=500
#create an array of the desired shape. 
resized = mx.zeros((new_h, new_w, c))
#Calculate width and height scaling factors
w_scale_factor = (old_w ) / (new_w )# if new_h != 0 else 0
h_scale_factor = (old_h ) / (new_h )# if new_w != 0 else 0
i_arr=mx.arange(0,new_h)
j_arr=mx.arange(0,new_w)

J_arr,I_arr=mx.meshgrid(j_arr,i_arr)
#map the coordinates back to the original image
x_arr=I_arr*h_scale_factor
y_arr=J_arr*w_scale_factor

#calculate the coordinate values for 4 surrounding pixels.
x_floor = mx.floor(x_arr).astype(mx.int32)
x_ceil = mx.minimum( old_h - 1, mx.ceil(x_arr)).astype(mx.int32) 
y_floor = mx.floor(y_arr).astype(mx.int32) 
y_ceil = mx.minimum(old_w - 1, mx.ceil(y_arr)).astype(mx.int32) 

x_arr_int=np.int32(x_arr)
y_arr_int=np.int32(y_arr)

# ### case 1: the pixel is an exact match
case1=(x_ceil==x_floor) & (y_ceil==y_floor)

resized[case1,:]= original_img[x_arr_int[case1], y_arr_int[case1], :]

[kyrollosyanny]

I believe, this is more of a limitation of mlx where all the methods have an output size same as its input size. I reposted this under a new issue. Thanks