commit 2f680b87418f88cefd9ac549c4c8577046227dd1 · pyrox.dev/nixpkgs

+2 -2

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

···

       50
        
             .${stdenv.hostPlatform.system};

     

       51
        
           hash =

     

       52
        
             {

     

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

     

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

     

       55
        
             }

     

       56
        
             .${stdenv.hostPlatform.system};

     

       57
        
         };

···

       50
        
             .${stdenv.hostPlatform.system};

     

       51
        
           hash =

     

       52
        
             {

     

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

     

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

     

       55
        
             }

     

       56
        
             .${stdenv.hostPlatform.system};

     

       57
        
         };

+6 -16

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

···

       36
        
       

     

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

     

       38
        
         srcs = {

     

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

     

       40
       -
             platform = "manylinux2014_x86_64";

     

       41
       -
             dist = "cp310";

     

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

     

       43
       -
           };

     

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

     

       45
       -
             platform = "manylinux2014_aarch64";

     

       46
       -
             dist = "cp310";

     

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

     

       48
       -
           };

     

       49
        
           "3.11-x86_64-linux" = getSrcFromPypi {

     

       50
        
             platform = "manylinux2014_x86_64";

     

       51
        
             dist = "cp311";

     

       52
       -
             hash = "sha256-CJbLswjZUpHiBc2J0lQCne46HfQ9ZumDEzGpr9LSeHA=";

     

       53
        
           };

     

       54
        
           "3.11-aarch64-linux" = getSrcFromPypi {

     

       55
        
             platform = "manylinux2014_aarch64";

     

       56
        
             dist = "cp311";

     

       57
       -
             hash = "sha256-LNjieaWaOLoMl4qDHhOt627p5Fcvujh8eXW6OtU13Tg=";

     

       58
        
           };

     

       59
        
           "3.12-x86_64-linux" = getSrcFromPypi {

     

       60
        
             platform = "manylinux2014_x86_64";

     

       61
        
             dist = "cp312";

     

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

     

       63
        
           };

     

       64
        
           "3.12-aarch64-linux" = getSrcFromPypi {

     

       65
        
             platform = "manylinux2014_aarch64";

     

       66
        
             dist = "cp312";

     

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

     

       68
        
           };

     

       69
        
           "3.13-x86_64-linux" = getSrcFromPypi {

     

       70
        
             platform = "manylinux2014_x86_64";

     

       71
        
             dist = "cp313";

     

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

     

       73
        
           };

     

       74
        
           "3.13-aarch64-linux" = getSrcFromPypi {

     

       75
        
             platform = "manylinux2014_aarch64";

     

       76
        
             dist = "cp313";

     

       77
       -
             hash = "sha256-dz76i1WoN0BsVh8O8CFE3akBkYEZN2DsVBnuyd0rmqw=";

     

       78
        
           };

     

       79
        
         };

     

       80
        
       in

···

       36
        
       

     

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

     

       38
        
         srcs = {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       39
        
           "3.11-x86_64-linux" = getSrcFromPypi {

     

       40
        
             platform = "manylinux2014_x86_64";

     

       41
        
             dist = "cp311";

     

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

     

       43
        
           };

     

       44
        
           "3.11-aarch64-linux" = getSrcFromPypi {

     

       45
        
             platform = "manylinux2014_aarch64";

     

       46
        
             dist = "cp311";

     

       47
       +
             hash = "sha256-JlpL8jjASPmJKeDa2XZCRtEkl73Yz98uIng6adc8xj4=";

     

       48
        
           };

     

       49
        
           "3.12-x86_64-linux" = getSrcFromPypi {

     

       50
        
             platform = "manylinux2014_x86_64";

     

       51
        
             dist = "cp312";

     

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

     

       53
        
           };

     

       54
        
           "3.12-aarch64-linux" = getSrcFromPypi {

     

       55
        
             platform = "manylinux2014_aarch64";

     

       56
        
             dist = "cp312";

     

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

     

       58
        
           };

     

       59
        
           "3.13-x86_64-linux" = getSrcFromPypi {

     

       60
        
             platform = "manylinux2014_x86_64";

     

       61
        
             dist = "cp313";

     

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

     

       63
        
           };

     

       64
        
           "3.13-aarch64-linux" = getSrcFromPypi {

     

       65
        
             platform = "manylinux2014_aarch64";

     

       66
        
             dist = "cp313";

     

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

     

       68
        
           };

     

       69
        
         };

     

       70
        
       in

+10 -2

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

···

       40
        
       in

     

       41
        
       buildPythonPackage rec {

     

       42
        
         pname = "jax";

     

       43
       -
         version = "0.6.2";

     

       44
        
         pyproject = true;

     

       45
        
       

     

       46
        
         src = fetchFromGitHub {

     
···

       48
        
           repo = "jax";

     

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

     

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

     

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

     

       52
        
         };

     

       53
        
       

     

       54
        
         build-system = [ setuptools ];

     
···

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

     

       104
        
           "tests/shape_poly_test.py::ShapePolyTest"

     

       105
        
           "tests/tree_util_test.py::TreeTest"

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       106
        
         ];

     

       107
        
       

     

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

···

       40
        
       in

     

       41
        
       buildPythonPackage rec {

     

       42
        
         pname = "jax";

     

       43
       +
         version = "0.7.0";

     

       44
        
         pyproject = true;

     

       45
        
       

     

       46
        
         src = fetchFromGitHub {

     
···

       48
        
           repo = "jax";

     

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

     

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

     

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

     

       52
        
         };

     

       53
        
       

     

       54
        
         build-system = [ setuptools ];

     
···

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

     

       104
        
           "tests/shape_poly_test.py::ShapePolyTest"

     

       105
        
           "tests/tree_util_test.py::TreeTest"

     

       106
       +
       

     

       107
       +
           # Mostly AssertionError on numerical tests failing since 0.7.0

     

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

     

       109
       +
           "tests/export_back_compat_test.py"

     

       110
       +
           "tests/lax_numpy_test.py"

     

       111
       +
           "tests/lax_scipy_test.py"

     

       112
       +
           "tests/lax_test.py"

     

       113
       +
           "tests/linalg_test.py"

     

       114
        
         ];

     

       115
        
       

     

       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
        
       }:

     

       19
        
       

     

       20
        
       let

     

       21
       -
         version = "0.6.2";

     

       22
        
         inherit (python) pythonVersion;

     

       23
        
       

     

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

     
···

       46
        
               };

     

       47
        
           in

     

       48
        
           {

     

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

     

       50
       -
               platform = "manylinux2014_x86_64";

     

       51
       -
               dist = "cp310";

     

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

     

       53
       -
             };

     

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

     

       55
       -
               platform = "manylinux2014_aarch64";

     

       56
       -
               dist = "cp310";

     

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

     

       58
       -
             };

     

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

     

       60
       -
               platform = "macosx_11_0_arm64";

     

       61
       -
               dist = "cp310";

     

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

     

       63
       -
             };

     

       64
       -
       

     

       65
        
             "3.11-x86_64-linux" = getSrcFromPypi {

     

       66
        
               platform = "manylinux2014_x86_64";

     

       67
        
               dist = "cp311";

     

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

     

       69
        
             };

     

       70
        
             "3.11-aarch64-linux" = getSrcFromPypi {

     

       71
        
               platform = "manylinux2014_aarch64";

     

       72
        
               dist = "cp311";

     

       73
       -
               hash = "sha256-Eern4FvFp5h12jYySvue3dS66u8qA4bK9tTzcgua7yg=";

     

       74
        
             };

     

       75
        
             "3.11-aarch64-darwin" = getSrcFromPypi {

     

       76
        
               platform = "macosx_11_0_arm64";

     

       77
        
               dist = "cp311";

     

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

     

       79
        
             };

     

       80
        
       

     

       81
        
             "3.12-x86_64-linux" = getSrcFromPypi {

     

       82
        
               platform = "manylinux2014_x86_64";

     

       83
        
               dist = "cp312";

     

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

     

       85
        
             };

     

       86
        
             "3.12-aarch64-linux" = getSrcFromPypi {

     

       87
        
               platform = "manylinux2014_aarch64";

     

       88
        
               dist = "cp312";

     

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

     

       90
        
             };

     

       91
        
             "3.12-aarch64-darwin" = getSrcFromPypi {

     

       92
        
               platform = "macosx_11_0_arm64";

     

       93
        
               dist = "cp312";

     

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

     

       95
        
             };

     

       96
        
       

     

       97
        
             "3.13-x86_64-linux" = getSrcFromPypi {

     

       98
        
               platform = "manylinux2014_x86_64";

     

       99
        
               dist = "cp313";

     

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

     

       101
        
             };

     

       102
        
             "3.13-aarch64-linux" = getSrcFromPypi {

     

       103
        
               platform = "manylinux2014_aarch64";

     

       104
        
               dist = "cp313";

     

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

     

       106
        
             };

     

       107
        
             "3.13-aarch64-darwin" = getSrcFromPypi {

     

       108
        
               platform = "macosx_11_0_arm64";

     

       109
        
               dist = "cp313";

     

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

     

       111
        
             };

     

       112
        
           };

     

       113
        
       in

···

       18
        
       }:

     

       19
        
       

     

       20
        
       let

     

       21
       +
         version = "0.7.0";

     

       22
        
         inherit (python) pythonVersion;

     

       23
        
       

     

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

     
···

       46
        
               };

     

       47
        
           in

     

       48
        
           {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       49
        
             "3.11-x86_64-linux" = getSrcFromPypi {

     

       50
        
               platform = "manylinux2014_x86_64";

     

       51
        
               dist = "cp311";

     

       52
       +
               hash = "sha256-B0oCVmTPQ5tZZdzKryDEqubMlV3d106FNCVoq6QN2kc=";

     

       53
        
             };

     

       54
        
             "3.11-aarch64-linux" = getSrcFromPypi {

     

       55
        
               platform = "manylinux2014_aarch64";

     

       56
        
               dist = "cp311";

     

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

     

       58
        
             };

     

       59
        
             "3.11-aarch64-darwin" = getSrcFromPypi {

     

       60
        
               platform = "macosx_11_0_arm64";

     

       61
        
               dist = "cp311";

     

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

     

       63
        
             };

     

       64
        
       

     

       65
        
             "3.12-x86_64-linux" = getSrcFromPypi {

     

       66
        
               platform = "manylinux2014_x86_64";

     

       67
        
               dist = "cp312";

     

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

     

       69
        
             };

     

       70
        
             "3.12-aarch64-linux" = getSrcFromPypi {

     

       71
        
               platform = "manylinux2014_aarch64";

     

       72
        
               dist = "cp312";

     

       73
       +
               hash = "sha256-AFITtty9ILC9ZVgLdAJVlQFQtInhpTBvZdjkn5EUq4U=";

     

       74
        
             };

     

       75
        
             "3.12-aarch64-darwin" = getSrcFromPypi {

     

       76
        
               platform = "macosx_11_0_arm64";

     

       77
        
               dist = "cp312";

     

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

     

       79
        
             };

     

       80
        
       

     

       81
        
             "3.13-x86_64-linux" = getSrcFromPypi {

     

       82
        
               platform = "manylinux2014_x86_64";

     

       83
        
               dist = "cp313";

     

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

     

       85
        
             };

     

       86
        
             "3.13-aarch64-linux" = getSrcFromPypi {

     

       87
        
               platform = "manylinux2014_aarch64";

     

       88
        
               dist = "cp313";

     

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

     

       90
        
             };

     

       91
        
             "3.13-aarch64-darwin" = getSrcFromPypi {

     

       92
        
               platform = "macosx_11_0_arm64";

     

       93
        
               dist = "cp313";

     

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

     

       95
        
             };

     

       96
        
           };

     

       97
        
       in

+1 -1

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

···

       10
        
           echo

     

       11
        
       }

     

       12
        
       

     

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

     

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

     

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

     

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

···

       10
        
           echo

     

       11
        
       }

     

       12
        
       

     

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

     

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

     

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

     

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