Jupyter notebook requires restarting kernel to make updated puzzle code effective

[CharlieTLe]

Thanks for creating this repo. I was going through the puzzles via the provided Jupyter notebook and noticed I had to restart the kernel to get the puzzle output to change.

To reproduce:

1. Set VERBOSE=1
2. Solve Puzzle 1

Generated source code for map:

[[kernel]] void custom_kernel_map(
  const constant int32_t* a [[buffer(0)]],
  device float* out [[buffer(1)]],
  uint3 thread_position_in_grid [[thread_position_in_grid]]) {

        uint local_i = thread_position_in_grid.x;
        // FILL ME IN (roughly 1 line)
        out[local_i] = a[local_i] + 10;
    
}

Passed Tests!

3. Update Puzzle to have an incorrect solution by adding 11 to each position instead of 10 and rerun the cells.

Generated source code for map:

[[kernel]] void custom_kernel_map(
  const constant int32_t* a [[buffer(0)]],
  device float* out [[buffer(1)]],
  uint3 thread_position_in_grid [[thread_position_in_grid]]) {

        uint local_i = thread_position_in_grid.x;
        [//](https://github.com/abeleinin/Metal-Puzzles/issues/new) FILL ME IN (roughly 1 line)
        out[local_i] = a[local_i] + 11;
    
}

Passed Tests!

Expected the test to fail, but the output still says that the tests passed. After restarting the kernel will the tests fail.

Generated source code for map:

[[kernel]] void custom_kernel_map(
  const constant int32_t* a [[buffer(0)]],
  device float* out [[buffer(1)]],
  uint3 thread_position_in_grid [[thread_position_in_grid]]) {

        uint local_i = thread_position_in_grid.x;
        [//](https://github.com/abeleinin/Metal-Puzzles/issues/new) FILL ME IN (roughly 1 line)
        out[local_i] = a[local_i] + 11;
    
}

Failed Tests.
Yours: array([11, 12, 13, 14], dtype=float32)
Spec : array([10, 11, 12, 13], dtype=int32)

Versions:

mlx==0.21.1

[abeleinin]

Thanks for submitting the issue! This is a side effect of the Metal JIT compiler, which stores functions in a static/global lookup table keyed by the function name. So, if you recompile something using the exact same name (ie custom_kernel_map), it'll simply returns that existing function pointer from the cache rather than actually recompiling. This isn't visible from the VERBOSE=1 output since I believe the function just writes using the source string during runtime.

The simple solution is to just rename the kernel using the name parameter in the MetalKernel class between executions. I found mx.metal.clear_cache(), but it doesn't seem to resolve the issue. Still looking into a longer term solution.