jax-test_cuda-passthru-failure

failure.log

ubuntu@bitbop:~/nixpkgs$ NIXPKGS_ALLOW_UNFREE=1 nixglhost -- nix run --impure .#python3Packages.jax.passthru.test_cuda_jaxlibBin
warning: Git tree '/home/ubuntu/nixpkgs' is dirty
Traceback (most recent call last):
  File "/nix/store/210z4yz9xpmch1fj7pzn83wdsbylmmlj-test_cuda/bin/test_cuda", line 7, in <module>
    rng = random.PRNGKey(0)
          ^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/random.py", line 240, in PRNGKey
    return _return_prng_keys(True, _key('PRNGKey', seed, impl))
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/random.py", line 202, in _key
    return prng.random_seed(seed, impl=impl)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/prng.py", line 595, in random_seed
    seeds_arr = jnp.asarray(np.int64(seeds))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py", line 2217, in asarray
    return array(a, dtype=dtype, copy=bool(copy), order=order)  # type: ignore
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py", line 2172, in array
    out_array: Array = lax_internal._convert_element_type(
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 560, in _convert_element_type
    return convert_element_type_p.bind(operand, new_dtype=new_dtype,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/core.py", line 444, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/core.py", line 447, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/core.py", line 935, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jax/_src/dispatch.py", line 87, in apply_primitive
    outs = fun(*args)
           ^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: NOT_FOUND: Couldn't find ptxas. The following locations were considered: /nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jaxlib/cuda/bin/ptxas, /home/ubuntu/.nix-profile/bin/ptxas, /nix/var/nix/profiles/default/bin/ptxas, /home/ubuntu/.vscode-server/bin/903b1e9d8990623e3d7da1df3d33db3e42d80eda/bin/remote-cli/ptxas, /home/ubuntu/.nix-profile/bin/ptxas, /nix/var/nix/profiles/default/bin/ptxas, /home/ubuntu/.nix-profile/bin/ptxas, /nix/var/nix/profiles/default/bin/ptxas, /usr/local/sbin/ptxas, /usr/local/bin/ptxas, /usr/sbin/ptxas, /usr/bin/ptxas, /sbin/ptxas, /bin/ptxas, /usr/games/ptxas, /usr/local/games/ptxas, /snap/bin/ptxas, /usr/local/cuda-11.8/bin/ptxas, /usr/local/cuda/bin/ptxas, /nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jaxlib/../nvidia/cuda_nvcc/bin/ptxas, /nix/store/w48pp30vl0jk2wvwdqgb9zg06my5m7sy-python3-3.11.7-env/lib/python3.11/site-packages/jaxlib/../../nvidia/cuda_nvcc/bin/ptxas
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

jax-default.nix

{ lib
, blas
, buildPythonPackage
, setuptools
, importlib-metadata
, fetchFromGitHub
, jaxlib
, jaxlib-bin
, hypothesis
, lapack
, matplotlib
, ml-dtypes
, numpy
, opt-einsum
, pkgs
, pytestCheckHook
, pytest-xdist
, pythonOlder
, scipy
, stdenv
}:

let
  usingMKL = blas.implementation == "mkl" || lapack.implementation == "mkl";
  # jaxlib is broken on aarch64-* as of 2023-03-05, but the binary wheels work
  # fine. jaxlib is only used in the checkPhase, so switching backends does not
  # impact package behavior. Get rid of this once jaxlib is fixed on aarch64-*.
  jaxlib' = if jaxlib.meta.broken then jaxlib-bin else jaxlib;

  jax = buildPythonPackage rec {
    pname = "jax";
    version = "0.4.24";
    pyproject = true;

    disabled = pythonOlder "3.9";

    src = fetchFromGitHub {
      owner = "google";
      repo = "jax";
      # google/jax contains tags for jax and jaxlib. Only use jax tags!
      rev = "refs/tags/${pname}-v${version}";
      hash = "sha256-hmx7eo3pephc6BQfoJ3U0QwWBWmhkAc+7S4QmW32qQs=";
    };

    nativeBuildInputs = [
      setuptools
    ];

    # The version is automatically set to ".dev" if this variable is not set.
    # https://github.com/google/jax/commit/e01f2617b85c5bdffc5ffb60b3d8d8ca9519a1f3
    JAX_RELEASE = "1";

    # jaxlib is _not_ included in propagatedBuildInputs because there are
    # different versions of jaxlib depending on the desired target hardware. The
    # JAX project ships separate wheels for CPU, GPU, and TPU.
    propagatedBuildInputs = [
      ml-dtypes
      numpy
      opt-einsum
      scipy
    ] ++ lib.optional (pythonOlder "3.10") importlib-metadata;

    nativeCheckInputs = [
      hypothesis
      jaxlib'
      matplotlib
      pytestCheckHook
      pytest-xdist
    ];

    # high parallelism will result in the tests getting stuck
    dontUsePytestXdist = true;

    # NOTE: Don't run the tests in the expiremental directory as they require flax
    # which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2.
    # Not a big deal, this is how the JAX docs suggest running the test suite
    # anyhow.
    pytestFlagsArray = [
      "--numprocesses=4"
      "-W ignore::DeprecationWarning"
      "tests/"
    ];

    disabledTests = [
      # Exceeds tolerance when the machine is busy
      "test_custom_linear_solve_aux"
      # UserWarning: Explicitly requested dtype <class 'numpy.float64'>
      #  requested in astype is not available, and will be truncated to
      # dtype float32. (With numpy 1.24)
      "testKde3"
      "testKde5"
      "testKde6"
      # Invokes python manually in a subprocess, which does not have the correct dependencies
      # ImportError: This version of jax requires jaxlib version >= 0.4.19.
      "test_no_log_spam"
    ] ++ lib.optionals usingMKL [
      # See
      #  * https://github.com/google/jax/issues/9705
      #  * https://discourse.nixos.org/t/getting-different-results-for-the-same-build-on-two-equally-configured-machines/17921
      #  * https://github.com/NixOS/nixpkgs/issues/161960
      "test_custom_linear_solve_cholesky"
      "test_custom_root_with_aux"
      "testEigvalsGrad_shape"
    ] ++ lib.optionals stdenv.isAarch64 [
      # See https://github.com/google/jax/issues/14793.
      "test_for_loop_fixpoint_correctly_identifies_loop_varying_residuals_unrolled_for_loop"
      "testQdwhWithRandomMatrix3"
      "testScanGrad_jit_scan"

      # See https://github.com/google/jax/issues/17867.
      "test_array"
      "test_async"
      "test_copy0"
      "test_device_put"
      "test_make_array_from_callback"
      "test_make_array_from_single_device_arrays"

      # Fails on some hardware due to some numerical error
      # See https://github.com/google/jax/issues/18535
      "testQdwhWithOnRankDeficientInput5"
    ];

    disabledTestPaths = lib.optionals (stdenv.isDarwin && stdenv.isAarch64) [
      # RuntimeWarning: invalid value encountered in cast
      "tests/lax_test.py"
    ];

    pythonImportsCheck = [ "jax" ];

    passthru =
      let
        test_cuda = jaxlib: pkgs.writers.writePython3Bin "test_cuda"
          {
            libraries = [ jax jaxlib ];
          } ''
          import jax
          from jax import random

          assert jax.devices()[0].platform == "gpu"

          rng = random.PRNGKey(0)
          x = random.normal(rng, (100, 100))
          x @ x

          print("success!")
        '';
      in
      {
        test_cuda_jaxlibSource = test_cuda (jaxlib.override { cudaSupport = true; });
        test_cuda_jaxlibBin = test_cuda (jaxlib-bin.override { cudaSupport = true; });
      };

    meta = with lib; {
      description = "Differentiate, compile, and transform Numpy code";
      homepage = "https://github.com/google/jax";
      license = licenses.asl20;
      maintainers = with maintainers; [ samuela ];
    };
  };
in
jax