commit eb491233cd74eb8a3642ca459ab672390ed02222 · pyrox.dev/nixpkgs

+6 -2

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

···

       27
       27
        
       in

     

       28
       28
        
       buildPythonPackage rec {

     

       29
       29
        
         pname = "jax";

     

       30
       30
       -
         version = "0.4.19";

     

       30
       30
       +
         version = "0.4.20";

     

       31
       31
        
         pyproject = true;

     

       32
       32
        
       

     

       33
       33
        
         disabled = pythonOlder "3.9";

     
···

       37
       37
        
           repo = "jax";

     

       38
       38
        
           # google/jax contains tags for jax and jaxlib. Only use jax tags!

     

       39
       39
        
           rev = "refs/tags/${pname}-v${version}";

     

       40
       40
       -
           hash = "sha256-l5uLPqhg/hqtO9oJSaioow5cH/0jKHDVziGezkfnVcc=";

     

       40
       40
       +
           hash = "sha256-WLYXUtchOaA6SGnKuVhN9CmV06xMCLQTEuEtL13ttZU=";

     

       41
       41
        
         };

     

       42
       42
        
       

     

       43
       43
        
         nativeBuildInputs = [

     
···

       108
       108
        
           "test_device_put"

     

       109
       109
        
           "test_make_array_from_callback"

     

       110
       110
        
           "test_make_array_from_single_device_arrays"

     

       111
       111
       +
       

     

       112
       112
       +
           # Fails on some hardware due to some numerical error

     

       113
       113
       +
           # See https://github.com/google/jax/issues/18535

     

       114
       114
       +
           "testQdwhWithOnRankDeficientInput5"

     

       111
       115
        
         ];

     

       112
       116
        
       

     

       113
       117
        
         disabledTestPaths = lib.optionals (stdenv.isDarwin && stdenv.isAarch64) [

+17 -17

pkgs/development/python-modules/jaxlib/bin.nix

···

       39
       39
        
       assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux;

     

       40
       40
        
       

     

       41
       41
        
       let

     

       42
       42
       -
         version = "0.4.19";

     

       42
       42
       +
         version = "0.4.20";

     

       43
       43
        
       

     

       44
       44
        
         inherit (python) pythonVersion;

     

       45
       45
        
       

     
···

       60
       60
        
             "3.9-x86_64-linux" = getSrcFromPypi {

     

       61
       61
        
               platform = "manylinux2014_x86_64";

     

       62
       62
        
               dist = "cp39";

     

       63
       63
       -
               hash = "sha256-8bTrWutuK0qVnbkcwMfgBf414YdaLc3GK5IsCm/JNPE=";

     

       63
       63
       +
               hash = "sha256-eIE+rz5x5BEkO85zncIWE8p/wDPxV8bnVJdHiknS998=";

     

       64
       64
        
             };

     

       65
       65
        
             "3.9-aarch64-darwin" = getSrcFromPypi {

     

       66
       66
        
               platform = "macosx_11_0_arm64";

     

       67
       67
        
               dist = "cp39";

     

       68
       68
       -
               hash = "sha256-Tmv2iOqlNbZqw/rYjef6GmM0N18EA5JTt6T3lQe+4Rs=";

     

       68
       68
       +
               hash = "sha256-dxInv8/aQiHsN7DpScuZao2ZyHDjF0AaTqUDA0qqg/M=";

     

       69
       69
        
             };

     

       70
       70
        
             "3.9-x86_64-darwin" = getSrcFromPypi {

     

       71
       71
        
               platform = "macosx_10_14_x86_64";

     

       72
       72
        
               dist = "cp39";

     

       73
       73
       -
               hash = "sha256-mDT1INLqPdCkxtMMFR0qHLOIZdWEy8Iuzw1/vOoECsA=";

     

       73
       73
       +
               hash = "sha256-wva6LkSokEHN+WQLCancVC7YBIxfImPsQpB1LzFcyqM=";

     

       74
       74
        
             };

     

       75
       75
        
       

     

       76
       76
        
             "3.10-x86_64-linux" = getSrcFromPypi {

     

       77
       77
        
               platform = "manylinux2014_x86_64";

     

       78
       78
        
               dist = "cp310";

     

       79
       79
       -
               hash = "sha256-ksnY+CPEstact5lKjbSg+ZSPJtSt0Y0NFWEFufBCByk=";

     

       79
       79
       +
               hash = "sha256-Yo2TYnkIelyy4vb5+nC/yY8SjV34i/jJvCe/VRQppmo=";

     

       80
       80
        
             };

     

       81
       81
        
             "3.10-aarch64-darwin" = getSrcFromPypi {

     

       82
       82
        
               platform = "macosx_11_0_arm64";

     

       83
       83
        
               dist = "cp310";

     

       84
       84
       -
               hash = "sha256-O7dHvdKLKfNELGfF4TKy7N5EX6Ca7Zu8OtLXWvFykR8=";

     

       84
       84
       +
               hash = "sha256-ufA/ACE4s4R/Fiq5SN7T44SVEN1Z5OfkJ/98lKxRFmo=";

     

       85
       85
        
             };

     

       86
       86
        
             "3.10-x86_64-darwin" = getSrcFromPypi {

     

       87
       87
        
               platform = "macosx_10_14_x86_64";

     

       88
       88
        
               dist = "cp310";

     

       89
       89
       -
               hash = "sha256-gqKMUZSXrt8sQtTAoQbzAfCzO8gM9Y1/tZpuJVWyN0Y=";

     

       89
       89
       +
               hash = "sha256-hBSrYQyOGMn0BexRWQKYnJdEYYlzHUWuWGHmjVT10TE=";

     

       90
       90
        
             };

     

       91
       91
        
       

     

       92
       92
        
             "3.11-x86_64-linux" = getSrcFromPypi {

     

       93
       93
        
               platform = "manylinux2014_x86_64";

     

       94
       94
        
               dist = "cp311";

     

       95
       95
       -
               hash = "sha256-m+NDzwXMNboNjDl2nLY+vqAoN2dQJZVWb1UQDpqqDPw=";

     

       95
       95
       +
               hash = "sha256-5N0nghTBrsa7d8kt8hZC2ghqlxCNC7U8ApD0PG7DHn8=";

     

       96
       96
        
             };

     

       97
       97
        
             "3.11-aarch64-darwin" = getSrcFromPypi {

     

       98
       98
        
               platform = "macosx_11_0_arm64";

     

       99
       99
        
               dist = "cp311";

     

       100
       100
       -
               hash = "sha256-zCOAjaWWCQT9Jnm1jjc1Rh5gemqy7ACtTKLM0MqSJzM=";

     

       100
       100
       +
               hash = "sha256-j13Br64cKe0hFh/cMBbOMuTXqauAvSKE+KzEmN7U6RA=";

     

       101
       101
        
             };

     

       102
       102
        
             "3.11-x86_64-darwin" = getSrcFromPypi {

     

       103
       103
        
               platform = "macosx_10_14_x86_64";

     

       104
       104
        
               dist = "cp311";

     

       105
       105
       -
               hash = "sha256-gOLIxkk+2hew2GqWu1WgMVEx1YEutx7Zod7QbwsuUVQ=";

     

       105
       105
       +
               hash = "sha256-nTnyawU4Ngq9VTE6oDuEfR6iJPRy+E/VLt98cU6eW4M=";

     

       106
       106
        
             };

     

       107
       107
        
       

     

       108
       108
        
             "3.12-x86_64-linux" = getSrcFromPypi {

     

       109
       109
        
               platform = "manylinux2014_x86_64";

     

       110
       110
        
               dist = "cp312";

     

       111
       111
       -
               hash = "sha256-BZTmkgNuV4nWtfbY4t/19aP43szZQEdgpFXh5qwGRXk=";

     

       111
       111
       +
               hash = "sha256-qPMoa7cso7DRBWuCJQoiOEzLPL3m76MPZZMYmZUj400=";

     

       112
       112
        
             };

     

       113
       113
        
             "3.12-aarch64-darwin" = getSrcFromPypi {

     

       114
       114
        
               platform = "macosx_11_0_arm64";

     

       115
       115
        
               dist = "cp312";

     

       116
       116
       -
               hash = "sha256-aAMTrLXU9EYwPv+kdeyI88/D7b4NANB39Fn8vuXUqFA=";

     

       116
       116
       +
               hash = "sha256-VqTC5egDHaDIvwVa3sAc9Sdtd0CwEFcXjDU/i54h844=";

     

       117
       117
        
             };

     

       118
       118
        
             "3.12-x86_64-darwin" = getSrcFromPypi {

     

       119
       119
        
               platform = "macosx_10_14_x86_64";

     

       120
       120
        
               dist = "cp312";

     

       121
       121
       -
               hash = "sha256-KHzlIfa9KtYcHX+i/F/SKaYTpD4/XjHVu5j3BdRTUmc=";

     

       121
       121
       +
               hash = "sha256-1F98Je2rMJJKrksI/EVAsX9n+dOpmDehUeAaMq/BY7o=";

     

       122
       122
        
             };

     

       123
       123
        
           };

     

       124
       124
        
       

     
···

       128
       128
        
         gpuSrcs = {

     

       129
       129
        
           "3.9" = fetchurl {

     

       130
       130
        
             url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl";

     

       131
       131
       -
             hash = "sha256-WB5Vbr/XeYKXCP/3DIXF20jR6/1xE3huX1h5ow8ETl0=";

     

       131
       131
       +
             hash = "sha256-VM2HuyMnG+hzrsTQEB5KJpqpBXyyp+eV1LVxmY1ZCGU=";

     

       132
       132
        
           };

     

       133
       133
        
           "3.10" = fetchurl {

     

       134
       134
        
             url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl";

     

       135
       135
       -
             hash = "sha256-zfN0n31+5GohwBkeQrqHus4qOyhM/GEdqG6KUupCZ4o=";

     

       135
       135
       +
             hash = "sha256-TLq3z3T2fjTcO3ESahboKG33mrOpjtj9C92f4d4nJKo=";

     

       136
       136
        
           };

     

       137
       137
        
           "3.11" = fetchurl {

     

       138
       138
        
             url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl";

     

       139
       139
       -
             hash = "sha256-Q8ZtF2GCrG30GFbCeCZTWPmW2TBybeXzh2u+NRiYpx4=";

     

       139
       139
       +
             hash = "sha256-CUXwyJq0HOo2j3Sw+NguBCnFkDuJpc3wfZUc90yyhOY=";

     

       140
       140
        
           };

     

       141
       141
        
           "3.12" = fetchurl {

     

       142
       142
        
             url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl";

     

       143
       143
       -
             hash = "sha256-lphkSDOJ9SwbO0hp/xC1bYn5fWgth9A9Iwsc9zV0buI=";

     

       143
       143
       +
             hash = "sha256-bAR8FLtiqufU+rL2a1q9c61CjH1eXxGTNGnDUkHlDBA=";

     

       144
       144
        
           };

     

       145
       145
        
         };

     

       146
       146

+6 -13

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

···

       54
       54
        
         inherit (cudaPackages) backendStdenv cudatoolkit cudaFlags cudnn nccl;

     

       55
       55
        
       

     

       56
       56
        
         pname = "jaxlib";

     

       57
       57
       -
         version = "0.4.19";

     

       57
       57
       +
         version = "0.4.20";

     

       58
       58
        
       

     

       59
       59
        
         meta = with lib; {

     

       60
       60
        
           description = "JAX is Autograd and XLA, brought together for high-performance machine learning research.";

     
···

       95
       95
        
           "absl_py"

     

       96
       96
        
           "astor_archive"

     

       97
       97
        
           "astunparse_archive"

     

       98
       98
       -
           "boringssl"

     

       99
       98
        
           # Not packaged in nixpkgs

     

       100
       99
        
           # "com_github_googleapis_googleapis"

     

       101
       100
        
           # "com_github_googlecloudplatform_google_cloud_cpp"

     
···

       151
       150
        
             repo = "jax";

     

       152
       151
        
             # google/jax contains tags for jax and jaxlib. Only use jaxlib tags!

     

       153
       152
        
             rev = "refs/tags/${pname}-v${version}";

     

       154
       154
       -
             hash = "sha256-l5uLPqhg/hqtO9oJSaioow5cH/0jKHDVziGezkfnVcc=";

     

       153
       153
       +
             hash = "sha256-WLYXUtchOaA6SGnKuVhN9CmV06xMCLQTEuEtL13ttZU=";

     

       155
       154
        
           };

     

       156
       155
        
       

     

       157
       156
        
           nativeBuildInputs = [

     
···

       264
       263
        
             ];

     

       265
       264
        
       

     

       266
       265
        
             sha256 = (if cudaSupport then {

     

       267
       267
       -
               x86_64-linux = "sha256-Hw4uFvltH7nlNN3qAEcQ+IR2FAOjRkvwyWA3rCPi7Vo=";

     

       266
       266
       +
               x86_64-linux = "sha256-QczClHxHElLZCqIZlHc3z3DXJ7rZQJaMs2XIb+lxarI=";

     

       268
       267
        
             } else {

     

       269
       269
       -
               x86_64-linux = "sha256-LEugnFwTV3EyeTZWgMvXzHbgeDPdmuT3daXCXJRMYVY=";

     

       270
       270
       -
               aarch64-linux = "sha256-59rv/3RjD8pnveBDZ33xZoNQxLmnhMocsKMgVfYoO70=";

     

       268
       268
       +
               x86_64-linux = "sha256-mqiJe4u0NYh1PKCbQfbo0U2e9/kYiBqj98d+BPHFSxQ=";

     

       269
       269
       +
               aarch64-linux = "sha256-EuLqamVBJ+qoVMCFIYUT846AghltZolfLGdtO9UeXSM=";

     

       271
       270
        
             }).${stdenv.system} or (throw "jaxlib: unsupported system: ${stdenv.system}");

     

       272
       271
        
           };

     

       273
       272
        
       

     
···

       293
       292
        
                 --replace "/usr/bin/install_name_tool" "${cctools}/bin/install_name_tool"

     

       294
       293
        
               substituteInPlace ../output/external/rules_cc/cc/private/toolchain/unix_cc_configure.bzl \

     

       295
       294
        
                 --replace "/usr/bin/libtool" "${cctools}/bin/libtool"

     

       296
       296
       -
             '' + (if stdenv.cc.isGNU then ''

     

       297
       297
       -
               sed -i 's@-lprotobuf@-l:libprotobuf.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

     

       298
       298
       -
               sed -i 's@-lprotoc@-l:libprotoc.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

     

       299
       299
       -
             '' else if stdenv.cc.isClang then ''

     

       300
       300
       -
               sed -i 's@-lprotobuf@${pkgs.protobuf}/lib/libprotobuf.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

     

       301
       301
       -
               sed -i 's@-lprotoc@${pkgs.protobuf}/lib/libprotoc.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

     

       302
       302
       -
             '' else throw "Unsupported stdenv.cc: ${stdenv.cc}");

     

       295
       295
       +
             '';

     

       303
       296
        
           };

     

       304
       297
        
       

     

       305
       298
        
           inherit meta;