python312Packages.flax: 0.10.3 -> 0.10.4

Diff: https://github.com/google/flax/compare/refs/tags/v0.10.3...v0.10.4

Changelog: https://github.com/google/flax/releases/tag/v0.10.4

pkgs/development/python-modules/flax/default.nix

@@ -39,14 +39,14 @@
 
 buildPythonPackage rec {
   pname = "flax";
-  version = "0.10.3";
+  version = "0.10.4";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "google";
     repo = "flax";
     tag = "v${version}";
-    hash = "sha256-PRKdtltiBVX9p6Sjw4sCDghqxYRxq4L9TLle1vy5dkk=";
+    hash = "sha256-+3PQPRVju9kw/4KWeifD8LhY4t6EzakhYISubMxrMw4=";
   };
 
   build-system = [
@@ -100,11 +100,55 @@ buildPythonPackage rec {
     "tests/jax_utils_test.py"
   ];
 
-  disabledTests = lib.optionals stdenv.hostPlatform.isDarwin [
-    # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
-    "test_ref_changed"
-    "test_structure_changed"
-  ];
+  disabledTests =
+    [
+      # AttributeError: module 'jax.api_util' has no attribute 'debug_info'
+      # https://github.com/google/flax/issues/4585
+      "test_basic_seq_lengths"
+      "test_bidirectional"
+      "test_big_resnet"
+      "test_custom_merge_fn"
+      "test_jit_scan_retracing_retracing"
+      "test_lazy_init"
+      "test_lazy_init"
+      "test_lazy_init_fails_on_data_dependence"
+      "test_lazy_init_fails_on_data_dependence"
+      "test_lifted_transform"
+      "test_lifted_transform_no_rename"
+      "test_multi_method_class_transform"
+      "test_numerical_equivalence"
+      "test_numerical_equivalence_single_batch"
+      "test_numerical_equivalence_single_batch_nn_scan"
+      "test_numerical_equivalence_with_mask"
+      "test_pjit_scan_over_layers"
+      "test_remat_scan"
+      "test_return_carry"
+      "test_reverse"
+      "test_reverse_but_keep_order"
+      "test_rnn_basic_forward"
+      "test_rnn_equivalence_with_flax_linen"
+      "test_rnn_multiple_batch_dims"
+      "test_rnn_time_major"
+      "test_rnn_unroll"
+      "test_rnn_with_spatial_dimensions"
+      "test_same_key"
+      "test_scan"
+      "test_scan_compact_count"
+      "test_scan_decorated"
+      "test_scan_negative_axes"
+      "test_scan_of_setup_parameter"
+      "test_scan_over_layers"
+      "test_scan_shared_params"
+      "test_scan_unshared_params"
+      "test_scan_with_axes"
+      "test_shared_cell"
+      "test_toplevel_submodule_adoption_pytree_transform"
+    ]
+    ++ lib.optionals stdenv.hostPlatform.isDarwin [
+      # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
+      "test_ref_changed"
+      "test_structure_changed"
+    ];
 
   passthru = {
     updateScript = writeScript "update.sh" ''

pkgs/development/python-modules/flaxlib/default.nix

@@ -1,40 +1,23 @@
 {
   lib,
-  stdenv,
   buildPythonPackage,
   flax,
   tomlq,
   python,
 
   # build-system
-  meson-python,
   nanobind,
   ninja,
+  scikit-build-core,
 
   # nativeBuildInputs
   cmake,
   pkg-config,
 }:
 
-let
-  nanobind-wrapper = stdenv.mkDerivation {
-    pname = "nanobind-wrapper";
-    inherit (nanobind) version;
-
-    src = ./nanobind-wrapper;
-
-    nativeBuildInputs = [
-      cmake
-    ];
-
-    buildFlags = [ "nanobind-static" ];
-
-    env.CMAKE_PREFIX_PATH = "${nanobind}/${python.sitePackages}/nanobind";
-  };
-in
 buildPythonPackage rec {
   pname = "flaxlib";
-  version = "0.0.1-a1";
+  version = "0.0.1";
   pyproject = true;
 
   inherit (flax) src;
@@ -43,11 +26,11 @@ buildPythonPackage rec {
 
   postPatch = ''
     expected_version="$version"
-    actual_version=$(${lib.getExe tomlq} --file Cargo.toml "package.version")
+    actual_version=$(${lib.getExe tomlq} --file pyproject.toml "project.version")
 
     if [ "$actual_version" != "$expected_version" ]; then
       echo -e "\n\tERROR:"
-      echo -e "\tThe version of the flaxlib python package ($expected_version) does not match the one in its Cargo.toml file ($actual_version)"
+      echo -e "\tThe version of the flaxlib python package ($expected_version) does not match the one in its pyproject.toml file ($actual_version)"
       echo -e "\tPlease update the version attribute of the nix python3Packages.flaxlib package."
       exit 1
     fi
@@ -56,15 +39,16 @@ buildPythonPackage rec {
   dontUseCmakeConfigure = true;
 
   build-system = [
-    meson-python
     nanobind
     ninja
+    scikit-build-core
   ];
   nativeBuildInputs = [
     cmake
     pkg-config
   ];
-  buildInputs = [ nanobind-wrapper ];
+
+  env.CMAKE_PREFIX_PATH = "${nanobind}/${python.sitePackages}/nanobind";
 
   pythonImportsCheck = [ "flaxlib" ];
 

pkgs/development/python-modules/flaxlib/nanobind-wrapper/CMakeLists.txt

@@ -1,9 +0,0 @@
-cmake_minimum_required(VERSION 3.31)
-project(nanobind-wrapper)
-
-find_package(Python COMPONENTS Interpreter Development REQUIRED)
-find_package(nanobind CONFIG REQUIRED)
-nanobind_build_library(nanobind-static)
-set_property(TARGET nanobind-static PROPERTY EXPORT_NAME nanobind)
-install(TARGETS nanobind-static EXPORT nanobind-static)
-install(EXPORT nanobind-static FILE nanobindConfig.cmake DESTINATION lib/nanobind/cmake)