tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 hatchling,
8
9 # dependencies
10 wadler-lindig,
11
12 # tests
13 cloudpickle,
14 equinox,
15 ipython,
16 jax,
17 jaxlib,
18 pytestCheckHook,
19 tensorflow,
20 torch,
21
22 # passthru
23 jaxtyping,
24}:
25
26buildPythonPackage (finalAttrs: {
27 pname = "jaxtyping";
28 version = "0.3.5";
29 pyproject = true;
30
31 src = fetchFromGitHub {
32 owner = "patrick-kidger";
33 repo = "jaxtyping";
34 tag = "v${finalAttrs.version}";
35 hash = "sha256-0Vt6UD1xQkwve6yDVi5XQCoJ/IsJWHCkGesj66myQq4=";
36 };
37
38 build-system = [ hatchling ];
39
40 dependencies = [
41 wadler-lindig
42 ];
43
44 pythonImportsCheck = [ "jaxtyping" ];
45
46 nativeCheckInputs = [
47 cloudpickle
48 equinox
49 ipython
50 jax
51 jaxlib
52 pytestCheckHook
53 tensorflow
54 torch
55 ];
56
57 doCheck = false;
58
59 # Enable tests via passthru to avoid cyclic dependency with equinox.
60 passthru.tests = {
61 check = jaxtyping.overridePythonAttrs {
62 # We disable tests because they complain about the version of typeguard being too new.
63 doCheck = false;
64 catchConflicts = false;
65 };
66 };
67
68 meta = {
69 description = "Type annotations and runtime checking for JAX arrays and PyTrees";
70 homepage = "https://github.com/patrick-kidger/jaxtyping";
71 changelog = "https://github.com/patrick-kidger/jaxtyping/releases/tag/${finalAttrs.src.tag}";
72 license = lib.licenses.mit;
73 maintainers = with lib.maintainers; [ GaetanLepage ];
74 };
75})