commit 2f680b87418f88cefd9ac549c4c8577046227dd1 · pyrox.dev/nixpkgs

+2 -2

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

···

       50
       50
        
             .${stdenv.hostPlatform.system};

     

       51
       51
        
           hash =

     

       52
       52
        
             {

     

       53
       53
       -
               x86_64-linux = "sha256-jNnq15SOosd4pQj+9dEVnot6v0/MxwN8P+Hb/NlQEtw=";

     

       54
       54
       -
               aarch64-linux = "sha256-IvrwINLo98oeKRVjMkH333Z4tzxwePXwsvETJIM3994=";

     

       53
       53
       +
               x86_64-linux = "sha256-dCzCYNaXdxy4gwhDmHSKAH0jMoaqP+X8FSp0wHPVmiQ=";

     

       54
       54
       +
               aarch64-linux = "sha256-3Oaxz+f0R0lCW66iKpAnQjuvsaUmnhPzB4XP0bLimGI=";

     

       55
       55
        
             }

     

       56
       56
        
             .${stdenv.hostPlatform.system};

     

       57
       57
        
         };

+6 -16

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

···

       36
       36
        
       

     

       37
       37
        
         # upstream does not distribute jax-cuda12-plugin 0.4.38 binaries for aarch64-linux

     

       38
       38
        
         srcs = {

     

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

     

       40
       40
       -
             platform = "manylinux2014_x86_64";

     

       41
       41
       -
             dist = "cp310";

     

       42
       42
       -
             hash = "sha256-F1H4iYkmmzzbDf5PewcqZEIUmBjJvJjDo5XIrK+RCnk=";

     

       43
       43
       -
           };

     

       44
       44
       -
           "3.10-aarch64-linux" = getSrcFromPypi {

     

       45
       45
       -
             platform = "manylinux2014_aarch64";

     

       46
       46
       -
             dist = "cp310";

     

       47
       47
       -
             hash = "sha256-vFw6ddBVGbTTJuRmnQ960P4PCs+HX5MT2RN0jMylqeo=";

     

       48
       48
       -
           };

     

       49
       39
        
           "3.11-x86_64-linux" = getSrcFromPypi {

     

       50
       40
        
             platform = "manylinux2014_x86_64";

     

       51
       41
        
             dist = "cp311";

     

       52
       52
       -
             hash = "sha256-CJbLswjZUpHiBc2J0lQCne46HfQ9ZumDEzGpr9LSeHA=";

     

       42
       42
       +
             hash = "sha256-2MmZTNYm4vFQmhZOXgfBIiIN3QptiDRvika9U8tePhk=";

     

       53
       43
        
           };

     

       54
       44
        
           "3.11-aarch64-linux" = getSrcFromPypi {

     

       55
       45
        
             platform = "manylinux2014_aarch64";

     

       56
       46
        
             dist = "cp311";

     

       57
       57
       -
             hash = "sha256-LNjieaWaOLoMl4qDHhOt627p5Fcvujh8eXW6OtU13Tg=";

     

       47
       47
       +
             hash = "sha256-JlpL8jjASPmJKeDa2XZCRtEkl73Yz98uIng6adc8xj4=";

     

       58
       48
        
           };

     

       59
       49
        
           "3.12-x86_64-linux" = getSrcFromPypi {

     

       60
       50
        
             platform = "manylinux2014_x86_64";

     

       61
       51
        
             dist = "cp312";

     

       62
       62
       -
             hash = "sha256-/r0Jn5cNNQ64+losmi+0sOp7PWqJ3xSWZj7fp6/lkOU=";

     

       52
       52
       +
             hash = "sha256-q37GQziSIY8lvdWi0TddqQh+mB36CYJ6x/NzzNu/LE0=";

     

       63
       53
        
           };

     

       64
       54
        
           "3.12-aarch64-linux" = getSrcFromPypi {

     

       65
       55
        
             platform = "manylinux2014_aarch64";

     

       66
       56
        
             dist = "cp312";

     

       67
       67
       -
             hash = "sha256-bJsALROx/LlANxPu3Th2oietH/vfs4EbH5+Jr0wlpfc=";

     

       57
       57
       +
             hash = "sha256-2UJ4q5a2ImQ3Ns3tt3eHA9MPL2T6SXrZ9bo13QyhBGQ=";

     

       68
       58
        
           };

     

       69
       59
        
           "3.13-x86_64-linux" = getSrcFromPypi {

     

       70
       60
        
             platform = "manylinux2014_x86_64";

     

       71
       61
        
             dist = "cp313";

     

       72
       72
       -
             hash = "sha256-20xhA8kS2M0a35TDTTE7tHYMp/AciXynzWLmXyeZQZk=";

     

       62
       62
       +
             hash = "sha256-v94NtxSyrV/p6GwfbS5gx9TZgNAVAphYIIgoq82CdM8=";

     

       73
       63
        
           };

     

       74
       64
        
           "3.13-aarch64-linux" = getSrcFromPypi {

     

       75
       65
        
             platform = "manylinux2014_aarch64";

     

       76
       66
        
             dist = "cp313";

     

       77
       77
       -
             hash = "sha256-dz76i1WoN0BsVh8O8CFE3akBkYEZN2DsVBnuyd0rmqw=";

     

       67
       67
       +
             hash = "sha256-5Nz+GFcW+0WsHti4kSIJbG2Tm0NRNxAzobz/JU3Rejg=";

     

       78
       68
        
           };

     

       79
       69
        
         };

     

       80
       70
        
       in

+10 -2

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

···

       40
       40
        
       in

     

       41
       41
        
       buildPythonPackage rec {

     

       42
       42
        
         pname = "jax";

     

       43
       43
       -
         version = "0.6.2";

     

       43
       43
       +
         version = "0.7.0";

     

       44
       44
        
         pyproject = true;

     

       45
       45
        
       

     

       46
       46
        
         src = fetchFromGitHub {

     
···

       48
       48
        
           repo = "jax";

     

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

     

       50
       50
        
           tag = "jax-v${version}";

     

       51
       51
       -
           hash = "sha256-MTgpwpJWxULCiZhDG+MFpOp8ZHoj1ZDmOD05OaGfXhM=";

     

       51
       51
       +
           hash = "sha256-CWtzT5c2Rzf2C0yaF9z2Y+uV3TdZC/h4AQf8t+U8o7o=";

     

       52
       52
        
         };

     

       53
       53
        
       

     

       54
       54
        
         build-system = [ setuptools ];

     
···

       103
       103
        
           "tests/pjit_test.py::PJitErrorTest::testAxisResourcesMismatch"

     

       104
       104
        
           "tests/shape_poly_test.py::ShapePolyTest"

     

       105
       105
        
           "tests/tree_util_test.py::TreeTest"

     

       106
       106
       +
       

     

       107
       107
       +
           # Mostly AssertionError on numerical tests failing since 0.7.0

     

       108
       108
       +
           # https://github.com/jax-ml/jax/issues/31428

     

       109
       109
       +
           "tests/export_back_compat_test.py"

     

       110
       110
       +
           "tests/lax_numpy_test.py"

     

       111
       111
       +
           "tests/lax_scipy_test.py"

     

       112
       112
       +
           "tests/lax_test.py"

     

       113
       113
       +
           "tests/linalg_test.py"

     

       106
       114
        
         ];

     

       107
       115
        
       

     

       108
       116
        
         # Prevents `tests/export_back_compat_test.py::CompatTest::test_*` tests from failing on darwin with

+10 -26

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

···

       18
       18
        
       }:

     

       19
       19
        
       

     

       20
       20
        
       let

     

       21
       21
       -
         version = "0.6.2";

     

       21
       21
       +
         version = "0.7.0";

     

       22
       22
        
         inherit (python) pythonVersion;

     

       23
       23
        
       

     

       24
       24
        
         # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the

     
···

       46
       46
        
               };

     

       47
       47
        
           in

     

       48
       48
        
           {

     

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

     

       50
       50
       -
               platform = "manylinux2014_x86_64";

     

       51
       51
       -
               dist = "cp310";

     

       52
       52
       -
               hash = "sha256-wIeg62+39vj1TVb0cwMo395QQN07Xd+oEOfCjqcQK0I=";

     

       53
       53
       -
             };

     

       54
       54
       -
             "3.10-aarch64-linux" = getSrcFromPypi {

     

       55
       55
       -
               platform = "manylinux2014_aarch64";

     

       56
       56
       -
               dist = "cp310";

     

       57
       57
       -
               hash = "sha256-QgXQmM6O+19/4v5QmLrmA2CU3I2IKfXg4NepsVUyYzY=";

     

       58
       58
       -
             };

     

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

     

       60
       60
       -
               platform = "macosx_11_0_arm64";

     

       61
       61
       -
               dist = "cp310";

     

       62
       62
       -
               hash = "sha256-2kYBsrXcjCPWr7KT6s+5rsTh0YccsvKcWhUdED5zsPg=";

     

       63
       63
       -
             };

     

       64
       64
       -
       

     

       65
       49
        
             "3.11-x86_64-linux" = getSrcFromPypi {

     

       66
       50
        
               platform = "manylinux2014_x86_64";

     

       67
       51
        
               dist = "cp311";

     

       68
       68
       -
               hash = "sha256-M11+NRXOeLUqQQE29Gqkp+oU0OfWQPNOHhN0CVVK0Kw=";

     

       52
       52
       +
               hash = "sha256-B0oCVmTPQ5tZZdzKryDEqubMlV3d106FNCVoq6QN2kc=";

     

       69
       53
        
             };

     

       70
       54
        
             "3.11-aarch64-linux" = getSrcFromPypi {

     

       71
       55
        
               platform = "manylinux2014_aarch64";

     

       72
       56
        
               dist = "cp311";

     

       73
       73
       -
               hash = "sha256-Eern4FvFp5h12jYySvue3dS66u8qA4bK9tTzcgua7yg=";

     

       57
       57
       +
               hash = "sha256-nfW6TIcSxVXs8y6yV0OR9kPlyg7KyiF4CEuMS/gktDM=";

     

       74
       58
        
             };

     

       75
       59
        
             "3.11-aarch64-darwin" = getSrcFromPypi {

     

       76
       60
        
               platform = "macosx_11_0_arm64";

     

       77
       61
        
               dist = "cp311";

     

       78
       78
       -
               hash = "sha256-ogj/YcWBKNMGu05a0IWL0rCWDywcEK1CxUj3SmDAAg4=";

     

       62
       62
       +
               hash = "sha256-/8tLHj4BIQb0OzBtcND2o2JigkoyT4n38ivyiGf76Bw=";

     

       79
       63
        
             };

     

       80
       64
        
       

     

       81
       65
        
             "3.12-x86_64-linux" = getSrcFromPypi {

     

       82
       66
        
               platform = "manylinux2014_x86_64";

     

       83
       67
        
               dist = "cp312";

     

       84
       84
       -
               hash = "sha256-8d0JtIGpPB1MdQAT9Gf3QZRJO6e9KfzU0c7BbjohT2U=";

     

       68
       68
       +
               hash = "sha256-U0+zJyuQ4sf47ZpCKaabXlwZsC+hRRbMxe750B8khUY=";

     

       85
       69
        
             };

     

       86
       70
        
             "3.12-aarch64-linux" = getSrcFromPypi {

     

       87
       71
        
               platform = "manylinux2014_aarch64";

     

       88
       72
        
               dist = "cp312";

     

       89
       89
       -
               hash = "sha256-h+wtycPtmrk27shTUWDF+9LISZSFWfHF2qdfY/q+WUI=";

     

       73
       73
       +
               hash = "sha256-AFITtty9ILC9ZVgLdAJVlQFQtInhpTBvZdjkn5EUq4U=";

     

       90
       74
        
             };

     

       91
       75
        
             "3.12-aarch64-darwin" = getSrcFromPypi {

     

       92
       76
        
               platform = "macosx_11_0_arm64";

     

       93
       77
        
               dist = "cp312";

     

       94
       94
       -
               hash = "sha256-NNimhKi+lJ3YfdSsyXEBtBBqDcmtFR7IkdoHIxmle5k=";

     

       78
       78
       +
               hash = "sha256-vpBV41zmvePpCfVbTLbttxR9CsLbCM+YbVw0EJhq+l0=";

     

       95
       79
        
             };

     

       96
       80
        
       

     

       97
       81
        
             "3.13-x86_64-linux" = getSrcFromPypi {

     

       98
       82
        
               platform = "manylinux2014_x86_64";

     

       99
       83
        
               dist = "cp313";

     

       100
       100
       -
               hash = "sha256-+UFj8UyP07qTrhS2Maus8UywMbugtZE4hpmEtNEDdfg=";

     

       84
       84
       +
               hash = "sha256-3zFmSlPBOpJjvKDow54DgKDMrgscElN232OkgNnLIIc=";

     

       101
       85
        
             };

     

       102
       86
        
             "3.13-aarch64-linux" = getSrcFromPypi {

     

       103
       87
        
               platform = "manylinux2014_aarch64";

     

       104
       88
        
               dist = "cp313";

     

       105
       105
       -
               hash = "sha256-cEmIN8r1OL1Fj/aFjIv9QE24IBWrqPZjZwGX+pkA/wI=";

     

       89
       89
       +
               hash = "sha256-W3OTyGlKF+1SLpVT4GeR3Xa0eJs0SNCF0O1P+613ouc=";

     

       106
       90
        
             };

     

       107
       91
        
             "3.13-aarch64-darwin" = getSrcFromPypi {

     

       108
       92
        
               platform = "macosx_11_0_arm64";

     

       109
       93
        
               dist = "cp313";

     

       110
       110
       -
               hash = "sha256-v/Z7GIEzzh8BEcexY6wyH9ZGtZ7SIepIkGPi4PhcuWc=";

     

       94
       94
       +
               hash = "sha256-Oo8ynwVNLggJPNWkr5MozOEsO1+rS9peLFza3GO17S0=";

     

       111
       95
        
             };

     

       112
       96
        
           };

     

       113
       97
        
       in

+1 -1

pkgs/development/python-modules/jaxlib/prefetch.sh

···

       10
       10
        
           echo

     

       11
       11
        
       }

     

       12
       12
        
       

     

       13
       13
       -
       for py in "310" "311" "312" "313"; do

     

       13
       13
       +
       for py in "311" "312" "313"; do

     

       14
       14
        
           prefetch "$py" "x86_64-linux" "jaxlib-bin"

     

       15
       15
        
           prefetch "$py" "aarch64-linux" "jaxlib-bin"

     

       16
       16
        
           prefetch "$py" "aarch64-darwin" "jaxlib-bin"