Commit 1616a488 authored by Fabian Affolter's avatar Fabian Affolter
Browse files

python311Packages.flax: refactor

parent 51ea18aa

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

+21 −14
Original line number Diff line number Diff line 
{ lib

, buildPythonPackage

, cloudpickle

, einops

, fetchFromGitHub

, jaxlib

, pythonRelaxDepsHook

, setuptools-scm

, jax

, jaxlib

, keras

, matplotlib

, msgpack

, numpy

, optax

, orbax-checkpoint

, pytest-xdist

, pytestCheckHook

, pythonOlder

, pythonRelaxDepsHook

, pyyaml

, rich

, setuptools-scm

, tensorflow

, tensorstore

, typing-extensions

, matplotlib

, cloudpickle

, einops

, keras

, pytest-xdist

, pytestCheckHook

, tensorflow

}:



buildPythonPackage rec {
@@ -26,6 +28,8 @@ buildPythonPackage rec { 
  version = "0.7.5";

  pyproject = true;



  disabled = pythonOlder "3.8";



  src = fetchFromGitHub {

    owner = "google";

    repo = "flax";
@@ -44,6 +48,7 @@ buildPythonPackage rec { 
    msgpack

    numpy

    optax

    orbax-checkpoint

    pyyaml

    rich

    tensorstore
@@ -75,7 +80,6 @@ buildPythonPackage rec { 
  disabledTestPaths = [

    # Docs test, needs extra deps + we're not interested in it.

    "docs/_ext/codediff_test.py"



    # The tests in `examples` are not designed to be executed from a single test

    # session and thus either have the modules that conflict with each other or

    # wrong import paths, depending on how they're invoked. Many tests also have
@@ -83,12 +87,15 @@ buildPythonPackage rec { 
    # `tensorflow_datasets`, `vocabulary`) so the benefits of trying to run them

    # would be limited anyway.

    "examples/*"



    # See https://github.com/google/flax/issues/3232.

    "tests/jax_utils_test.py"

    # Requires tree

    "tests/tensorboard_test.py"

  ];



    # Requires orbax which is not packaged as of 2023-07-27.

    "tests/checkpoints_test.py"

  disabledTests = [

    # ValueError: Checkpoint path should be absolute

    "test_overwrite_checkpoints0"

  ];



  meta = with lib; {