Maybe theres a better way, but I couldn't find a satisfactory answer online so heres the code i use if someone needs it: import numpy as np def np_gather(array: np.ndarray, indices: np.ndarray, axis: tuple, value: np.ndarray = None): """ gather and affect indices behaviour along different non contiguous axis """ assert isinstance(axis, tuple) assert len(axis) > 0 shape_o = array.shape axis_o = np.arange(array.ndim) permute_shape = axis_o[np.logical_not(np.isin(axis_o, axis))].tolist() permute_shape.extend(axis) safe_dims = np.array(shape_o)[np.logical_not(np.isin(axis_o, axis))] unsafe_dims = np.array(shape_o)[np.isin(axis_o, axis)] array = np.transpose(array, permute_shape) array = array.reshape((*safe_dims, -1)) offsets = np.cumprod(np.insert(unsafe_dims, unsafe_dims.shape[0], 1)[1:][::-1])[::-1] indices_temp = (indices * offsets.reshape(indices.shape[0], 1)).sum(axis=0) if not value is None: array[..., indices_temp] = value.reshape(*[1 for d in safe_dims], *value.shape) return array[..., indices_temp] Example : array1 = np.random.randn(10,10) check = array1>0 array2 = np.tile(array1.reshape(1,1,10,10),(4,6,1,1)).transpose(0,2,1,3) nix = np.stack(np.nonzero(check)) print((np_gather(array2,nix,(1,3))>0).all()) new_value = np.ones(nix.shape[-1])*0.12 returned_values = np_gather(array2,nix,(1,3),new_value) print((returned_values==0.12).all()) Also it's not letting me post the code because there is not enough text, so here is some, blabla Continue reading...
