Added lock_shard guard to keep_alive

[vfdev-5]

Hi @wjakob in this PR I added lock_shard guard(shard); which may be missed in your free-threading support PR to keep_alive methods as I saw everywhere this pattern.

I also have a TSAN report talking about this place when running an extension with class with nb::is_weak_referenceable.

WARNING: ThreadSanitizer: data race (pid=213447)
  Read of size 8 at 0x7b8c000011d0 by thread T1:
    #0 tsl::rh::power_of_two_growth_policy<2ul>::bucket_for_hash(unsigned long) const /usr/local/lib/python3.13t/dist-packages/nanobind/ext/robin_map/include/tsl/robin_growth_policy.h:125 (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x45dda)
    #1 tsl::detail_robin_hash::robin_hash<std::pair<void*, void*>, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::KeySelect, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::bucket_for_hash(unsigned long) const /usr/local/lib/python3.13t/dist-packages/nanobind/ext/robin_map/include/tsl/robin_hash.h:1133 (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x546c1)
    #2 std::pair<tsl::detail_robin_hash::robin_hash<std::pair<void*, void*>, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::KeySelect, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::robin_iterator<false>, bool> tsl::detail_robin_hash::robin_hash<std::pair<void*, void*>, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::KeySelect, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::insert_impl<void*, std::piecewise_construct_t const&, std::tuple<void*&&>, std::tuple<> >(void* const&, std::piecewise_construct_t const&, std::tuple<void*&&>&&, std::tuple<>&&) /usr/local/lib/python3.13t/dist-packages/nanobind/ext/robin_map/include/tsl/robin_hash.h:1222 (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x56873)
    #3 std::pair<tsl::detail_robin_hash::robin_hash<std::pair<void*, void*>, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::KeySelect, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::robin_iterator<false>, bool> tsl::detail_robin_hash::robin_hash<std::pair<void*, void*>, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::KeySelect, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::try_emplace<void*>(void*&&) /usr/local/lib/python3.13t/dist-packages/nanobind/ext/robin_map/include/tsl/robin_hash.h:809 (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x54569)
    #4 tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect::value_type& tsl::detail_robin_hash::robin_hash<std::pair<void*, void*>, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::KeySelect, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::operator[]<void*, tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::ValueSelect, (void*)0>(void*&&) <null> (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x52a87)
    #5 tsl::robin_map<void*, void*, nanobind::detail::ptr_hash, std::equal_to<void*>, std::allocator<std::pair<void*, void*> >, false, tsl::rh::power_of_two_growth_policy<2ul> >::operator[](void*&&) /usr/local/lib/python3.13t/dist-packages/nanobind/ext/robin_map/include/tsl/robin_map.h:457 (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x51e5b)
    #6 nanobind::detail::keep_alive(_object*, void*, void (*)(void*) noexcept) /usr/local/lib/python3.13t/dist-packages/nanobind/src/nb_type.cpp:1612 (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x4daad)
    #7 nanobind::class_<jax::WeakrefLRUCache>::class_<nanobind::is_weak_referenceable>(nanobind::handle, char const*, nanobind::is_weak_referenceable const&)::{lambda(_object*)#1}::operator()(_object*) const <null> (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x222fd)
    #8 nanobind::class_<jax::WeakrefLRUCache>::class_<nanobind::is_weak_referenceable>(nanobind::handle, char const*, nanobind::is_weak_referenceable const&)::{lambda(_object*)#1}::_FUN(_object*) <null> (wrlru_cache_ext.cpython-313t-x86_64-linux-gnu.so+0x22412)

I rerun TSAN with this PR and it is happy now, however, feel free to close this PR if it was an intended omission and please suggest how to fix TSAN.
Thanks a lot!

[wjakob]

Indeed, I missed that one. Thank you!

[wjakob]

(and yes, this is the right place to put the lock)