jax

2c52ae35 - Fixed deadlock in MakeShardFn on static var assignment under free-threading

Fixed deadlock in MakeShardFn on static var assignment under free-threading Description: - Fixed deadlock in MakeShardFn on static var assignment under free-threading using nb::ft_mutex Fixes #28385