Unverified Commit 2da53957 authored by Gaétan Lepage's avatar Gaétan Lepage Committed by GitHub
Browse files

python3Packages.distrax: 0.1.5 -> 0.1.7 (#439370)

parents c7d251be 4f613e09

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

+15 −25
Original line number Diff line number Diff line
@@ -2,11 +2,19 @@ 
  lib,

  buildPythonPackage,

  fetchFromGitHub,

  fetchpatch2,



  # build-system

  flit-core,



  # dependencies

  absl-py,

  chex,

  jax,

  jaxlib,

  numpy,

  tensorflow-probability,



  # tests

  dm-haiku,

  pytest-xdist,

  pytestCheckHook,
@@ -14,42 +22,24 @@ 


buildPythonPackage rec {

  pname = "distrax";

  version = "0.1.5";

  version = "0.1.7";

  pyproject = true;



  src = fetchFromGitHub {

    owner = "google-deepmind";

    repo = "distrax";

    tag = "v${version}";

    hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0=";

    hash = "sha256-R6rGGNzup3O6eZ2z4vygYWTjroE/Irt3aog8Op+0hco=";

  };



  patches = [

    # TODO: remove at the next release (already on master)

    (fetchpatch2 {

      name = "fix-jax-0.6.0-compat";

      url = "https://github.com/google-deepmind/distrax/commit/c02708ac46518fac00ab2945311e0f2ee32c672c.patch";

      hash = "sha256-hFNXKoA1b5I6dzhwTRXp/SnkHv89GI6tYwlnBBHwG78=";

    })

    # https://github.com/google-deepmind/distrax/pull/289

    (fetchpatch2 {

      name = "fix-jax-0.7.0-compat";

      url = "https://github.com/google-deepmind/distrax/commit/7fc5bd7efff4a7144d175199159f115c3e68a3cf.patch";

      hash = "sha256-TiD72YIb6ajpaCO1yOGl/+JCuaikQ879Zcpaf2wzMq4=";

    })

  build-system = [

    flit-core

  ];



  # TODO: remove at the next release (already on master)

  # https://github.com/google-deepmind/distrax/pull/293

  postPatch = ''

    substituteInPlace distrax/_src/utils/transformations.py \

      --replace-fail \

        "jax.experimental.pjit.pjit_p" \

        "jex.core.primitives.jit_p"

  '';



  dependencies = [

    absl-py

    chex

    jax

    jaxlib

    numpy

    tensorflow-probability