python310Packages.chex: 0.1.6 -> 0.1.82; unbreak

{ absl-py

{ lib

, buildPythonPackage

, pythonOlder

, fetchFromGitHub

, absl-py

, cloudpickle

, dm-tree

, fetchFromGitHub

, jax

, jaxlib

, lib

, numpy

, pytestCheckHook

, toolz

, typing-extensions

}:

buildPythonPackage rec {

  pname = "chex";

  version = "0.1.6";

  version = "0.1.82";

  format = "setuptools";

  disabled = pythonOlder "3.9";

  src = fetchFromGitHub {

    owner = "deepmind";

    repo = pname;

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

    hash = "sha256-VolRlLLgKga9S17ByVrYya9VPtu9yiOnvt/WmlE1DOc=";

    hash = "sha256-xBq22AaR2Tp1NSPefEyvCDeUYqRZlAf5LVHWo0luiXk=";

  };

  propagatedBuildInputs = [

    absl-py

    cloudpickle

    dm-tree

    jaxlib

    jax

    numpy

    toolz

    typing-extensions

  ];

  pythonImportsCheck = [

  ];

  nativeCheckInputs = [

    jaxlib

    cloudpickle

    dm-tree

    pytestCheckHook

  ];

  disabledTests = [

    # See https://github.com/deepmind/chex/issues/204.

    "test_uninspected_checks"

    # These tests started failing at some point after upgrading to 0.1.5

    "test_useful_failure"

    "TreeAssertionsTest"

    "PmapFakeTest"

    "WithDeviceTest"

  ];

  meta = with lib; {

    description = "Chex is a library of utilities for helping to write reliable JAX code.";

    homepage = "https://github.com/deepmind/chex";