1{ lib
2, buildPythonPackage
3, fetchPypi
4
5# dependencies
6, torch
7, scipy
8
9# tests
10, pytestCheckHook
11}:
12
13buildPythonPackage rec {
14 pname = "torchdiffeq";
15 version = "0.2.3";
16 format = "setuptools";
17
18 src = fetchPypi {
19 inherit pname version;
20 hash = "sha256-/nX0NLkJCsDCdwLgK+0hRysPhwNb5lgfUe3F1AE+oxo=";
21 };
22
23 propagatedBuildInputs = [
24 torch
25 scipy
26 ];
27
28 pythonImportsCheck = [ "torchdiffeq" ];
29
30 # no tests in sdist, no tags on git
31 doCheck = false;
32
33 nativeCheckInputs = [
34 pytestCheckHook
35 ];
36
37 meta = with lib; {
38 description = "Differentiable ODE solvers with full GPU support and O(1)-memory backpropagation";
39 homepage = "https://github.com/rtqichen/torchdiffeq";
40 license = licenses.mit;
41 maintainers = teams.tts.members;
42 };
43}