Custom Kernel Extra Documentation

[ppoage]

I've been playing around with the custom kernels to try and get around some functions that aren't implemented from numpy. (or in this case scipy peak detection). Mainly because you don't yet support include numpy.nonzero() and the single input version of numpy.where()

I'm trying to make a np.nonzero function to go in something similar to this:

# want to make _mlx version for mac
def _find_peaks_default(x, height=None, threshold=None, distance=None,
               prominence=None, width=None, wlen=None, rel_height=0.5,
               plateau_size=None):
    """
    NumPy version of scipy.signal.find_peaks.
    Returns (peaks, properties) where:
      peaks      = 1D array of peak indices
      properties = dict of requested peak properties
    """
    
    ...
    
    # local maxima mask
    mask = np.empty(N, bool)
    mask[0]  = mask[-1] = False
    mask[1:-1] = (x[1:-1] > x[:-2]) & (x[1:-1] > x[2:])
    peaks = np.nonzero(mask)[0]

I'm trying to find what other documentation there is on the custom kernel. Namely, output shape. I'd love to be able to do something like the given example

def grid_sample(x, grid):

...

@grid_sample.vjp
def grid_sample_vjp(primals, cotangent, _):

except with the ability to infer the second kernel's output shape if the first calculates it. Otherwise I would have to assign a fixed buffer and then cut it to size. Are indirect command buffers supported? I'm new to MSL but it looks like it could possibly do stuff like this.

[awni]

In general you can't do data-dependent output shapes in MLX. You can read more about why here.

Assigning a fixed buffer and cutting it is probably your best bet.

Are indirect command buffers supported?

No sorry