diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index 57454d4f84b6..9b640d02fbcb 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -50,8 +50,8 @@ buildPythonPackage (finalAttrs: { .${stdenv.hostPlatform.system}; hash = { - x86_64-linux = "sha256-MJUVRT9Zyq2Vv3bIvGScJLwOPRLQe68895K+CCq97js="; - aarch64-linux = "sha256-NqwxbFiKDgjrZPhoJwUmrco4VtS3CxN2C36bc76F0s8="; + x86_64-linux = "sha256-gG0f0pA4tqz1orKJ3WIZKr6pd+4mrvYOopWx0oojrPg="; + aarch64-linux = "sha256-tvLWZUjuHukQqDa1/j0Ou+HaBKYtD0aLpIIhiQNqMrM="; } .${stdenv.hostPlatform.system}; }; diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix index ad450298399c..01eb6ec63239 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -39,42 +39,42 @@ let "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp311"; - hash = "sha256-lXMXyv8eawBqE1UXM3cN8jWidHBkviawJ0YP46WPWO4="; + hash = "sha256-drv+zi9u/7TokI64pfvxq7wRJd5mdoPsL/ulHGPcV68="; }; "3.11-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp311"; - hash = "sha256-GkpKVqgyrTDgHBQIWx+iRApV+kzggI5AXnNlDv+K+zE="; + hash = "sha256-nNCTcHC70o5EuRQ1tSxyJcjtbLN922qQ6cJGhzoD+c0="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; - hash = "sha256-oWugDXI2YUWCfKOajISfn+4Rd9ZzjJ7f1ys6n+31zxg="; + hash = "sha256-Trbo4Jkv2Yl9skaNoz8nauQU7ITcQ2gEEkQEQwUC7qw="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; - hash = "sha256-6vGmGVqg60i0IMcfIrktGXeNlIAkl3yuRkEnnV9sH4E="; + hash = "sha256-dnoUgtv2UmiEA8TiL/AUcOGt0TxI+dgXFp/Y90QK/OU="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; - hash = "sha256-7QdNN5Z7zmHBuv3fndCoMCMGwS+QcAylrfo9CP3Uruc="; + hash = "sha256-lIqYiSfLEIQ7UBohUYHMsdrtP6cFMbLxrghC/hwIsF4="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; - hash = "sha256-dHnmfZ2Y1KD1JflPi4GIDcwFhjH23hGLkd6R2xtAle0="; + hash = "sha256-aXdqyEkRL0/C1vphb4dyEG2v84/NuCXE459H6HjydhA="; }; "3.14-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp314"; - hash = "sha256-96ahiBNJ3DnyRIGioBn/KUuUkJuAqIge/iHUW2pfaR0="; + hash = "sha256-DHoCBBVVhcwcT8swxmt/iBs45X4fDY6X1I4WVL8NJ/g="; }; "3.14-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp314"; - hash = "sha256-iIb8V5df7gG+FyX4mLaSkvYVtwgwrvXC+95jgzNha10="; + hash = "sha256-SUbsV923leLagojOsc7d1YwqIEoPd2WvierkhqHTrOc="; }; }; in diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index a8a84ae2747b..b1e084faa030 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -6,7 +6,6 @@ lapack, buildPythonPackage, fetchFromGitHub, - fetchpatch, cudaSupport ? config.cudaSupport, # build-system @@ -41,7 +40,7 @@ let in buildPythonPackage (finalAttrs: { pname = "jax"; - version = "0.10.0"; + version = "0.10.2"; pyproject = true; __structuredAttrs = true; @@ -50,18 +49,9 @@ buildPythonPackage (finalAttrs: { repo = "jax"; # google/jax contains tags for jax and jaxlib. Only use jax tags! tag = "jax-v${finalAttrs.version}"; - hash = "sha256-/RCihrjONN/+QwyQRNEmlIa7JsCLzz+SkBe5sd+ThgU="; + hash = "sha256-OQkh9uC8NsxoG3SByPybXQ81c11T3lYgjaU3tbB0+6E="; }; - patches = [ - # setup.py: Include only jax.* in the built wheel - # https://github.com/jax-ml/jax/pull/37182 - (fetchpatch { - url = "https://github.com/jax-ml/jax/commit/cb5a91780f84f124090d8f94e99c8771e87590f6.patch"; - hash = "sha256-p6X9IFe4YUb2MQp7YjJHme1dueZ1Y37IKnANGruW1cM="; - }) - ]; - build-system = [ setuptools ]; # The version is automatically set to ".dev" if this variable is not set. @@ -135,6 +125,13 @@ buildPythonPackage (finalAttrs: { disabledTests = [ # Exceeds tolerance when the machine is busy "test_custom_linear_solve_aux" + + # pytest-xdist/execnet cannot serialize the numpy `type` objects this test passes to + # self.subTest(dtype=...) when shipping subtest reports between workers. + # The assertions themselves pass; the failure is a harness artifact of running with + # --numprocesses. + # New test in jax 0.10.2 (tests/random_impl_test.py). + "test_random_bits" ] ++ lib.optionals usingMKL [ # See diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index deb80c9b68ad..1dce7d1025ad 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -18,7 +18,7 @@ }: let - version = "0.10.0"; + version = "0.10.2"; inherit (python) pythonVersion; # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the @@ -49,65 +49,65 @@ let "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp311"; - hash = "sha256-m+IpmTpB5bK4TyNOzBml3gLzXt2xGVzwJ71TnhYB4V0="; + hash = "sha256-H6yjxdRmLLSmEwpoEF1ou1IHZIF+Fl1u6/1nhsDR8w8="; }; "3.11-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp311"; - hash = "sha256-PblOvIWTddlV3jUEGCrdfOFzPOPTDBXg7wMWAstRpVk="; + hash = "sha256-1EVl3P0bT2D3bZEcZRIRiopPx2S975JmP+y4v8zVTyM="; }; "3.11-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp311"; - hash = "sha256-J3Ay6fB0w/1f/R4MsD1P5m4nLeRyZnzbxBitmbIbZGo="; + hash = "sha256-WpiHP8hnYjuB8r7hXVVLjt1liKGD0B+lDSGx49uW/ys="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; - hash = "sha256-sL+4ZaB98ubXQYwLDCkt0pS1UAUjsd1YcrGA2yqkgNQ="; + hash = "sha256-/ojsRDcUxDeZaLbBCfn6YXx60ZuAKCjk17+GHNZtpLc="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; - hash = "sha256-qh1w8aTifrQDZU5x4vso1XhtPpt3/BhH6MU4mICSfKQ="; + hash = "sha256-U7cpd65YLAOp6OHN7h77+OvBQYJwllsOaereV6z0AzE="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; - hash = "sha256-fB2bRjMnx6IzPyEBFOywTyj+/FG6gjOoWiKAzOdb20I="; + hash = "sha256-R7t8ARUV6oYr5+gxP0D5xWy+wJ3Jig/LUBZ4X81FTAE="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; - hash = "sha256-0wPcMbZei3k9VgD4GxWDvgPcm4dqTBCz4lm2YJocvjs="; + hash = "sha256-nkgYtKh1b9ORh2bKKqU0ISWAn08Ipv5GAm1DhufCNkQ="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; - hash = "sha256-bY14twcLNOTFu6X34Qkn5/Sqybab4X6bCliYVTpDOPM="; + hash = "sha256-RbKLAjhperdLvPIEEar7bbQqzDGDbML9cR5c8Fa/lVY="; }; "3.13-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp313"; - hash = "sha256-OEY1//VYmaKVu8gu5sb3c6MA54fcRyypK755q/qsg2k="; + hash = "sha256-TfUwr6NUoi3BdHpdVgZARQy7iV1JiJM4o/WMdqTHbI4="; }; "3.14-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp314"; - hash = "sha256-KkLPBMD4i8A7FQoX+n3bsvQOCWZn7IobhA7YeRPm5zU="; + hash = "sha256-yjTzYxl/sKxAglgsp1UAeRA2njP4qLo9Ne2UtxBwEH0="; }; "3.14-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp314"; - hash = "sha256-rUfgckMJeewhY3qkh9TcRkAouOm+JyaPN95pU2x240E="; + hash = "sha256-8Y9W/ukGmc+6m2YnBFp6KZcCyw4q+CzhgNmmp8gEgJM="; }; "3.14-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp314"; - hash = "sha256-mLJmcpQ2cnQoc/ZbwDIWgZ/FUyXJnxRlkNAHwBcr/zA="; + hash = "sha256-cuuiixL+4CYW+kKqS4gbSrYtd1fHhDxGJAHT+zSie+Q="; }; }; in