[andresy]

No description provided.

[davidkoski]

To make sure I understand, you use this like:

uintptr_t id = ...;
mlx_compile_cache cache;

// this is roughly attaching the cache to the global?
mlx_detail_compile_cache(&mlx_detail_compile_cache);

mlx_detail_compile_erase(cache, id);

There is no lifecycle on cache? I am not sure what the purpose of the variable is since it seems to be picking up a global.

[zcbenz]

mlx_detail_compile_cache returns the cache for current thread and erase/clear_cache takes it to clear the cache on the correct thread, this was used to fix a python bug that a compiled function could be deleted on a thread different from where it was compiled.

[davidkoski]

Ah, interesting. In swift we may allocate on one thread and deallocate on another thread -- indeed the thread may be gone by the time we delete. Ownership of the compiled function (the closure we retain) is thread safe on the swift side. Can we call mlx_detail_compile_erase() from any thread?

This may be along the lines of the stream: threads are not necessarily the unit of execution in swift.

[zcbenz]

Yeah mlx_detail_compile_erase() is thread safe to call from anywhere.

Also we do need another public API to free mlx_compile_cache, it is a C++ weak ptr.

[davidkoski]

This would close #117

[davidkoski]

Should this be an mlx_string?

[davidkoski]

I think this is missing overloads:

  m.def(
      "count_nonzero",
      [](const mx::array& a,
         const IntOrVec& axis,
         bool keepdims,
         mx::StreamOrDevice s) {
        if (std::holds_alternative<std::monostate>(axis)) {
          return mx::count_nonzero(a, keepdims, s);
        } else if (auto pv = std::get_if<int>(&axis); pv) {
          return mx::count_nonzero(a, *pv, keepdims, s);
        } else {
          return mx::count_nonzero(
              a, std::get<std::vector<int>>(axis), keepdims, s);
        }
      },