Jupyter caching shaders

[hlovatt]

Hi,

I'm writing shaders using mx.fast.metal_kernal in Jupyter. When I change the shader code it doesn't update without first quitting the Jupyter Kernel (not to be confused with the Metal Kernel!). The Metal Kernel must be cached (or similar) somewhere.

Does anyone know how to delete the old Metal Kernel and replace it with the new one in Jupyter?

EG

import mlx.core as mx
def exp_elementwise(a: mx.array):
    source = '''
        uint elem = thread_position_in_grid.x;
        out[elem] = 0.0;
        //out[elem] = inp[elem];
        out[elem] = metal::exp(out[elem]);
    '''
    kernel = mx.fast.metal_kernel(
        name="myexp",
        input_names=["inp"],
        output_names=["out"],
        source=source
    )
    outputs = kernel(
        inputs=[a],
        template=[("T", a.dtype)],
        grid=(a.size, 1, 1),
        threadgroup=(256, 1, 1),
        output_shapes=[a.shape],
        output_dtypes=[a.dtype],
        verbose=True,
    )
    return outputs[0]

ax = mx.random.normal(shape=(4, 16)).astype(mx.float16)
exp_elementwise(ax)

Run above and it will produce an array of all 1's - as expected. But, uncomment out[elem] = input[elem] and run again without quitting the Jupyter Kernel and you get the same result!