Aleaverfay/temp sparse dispatch fa elec

.buildkite/bin/linting

@@ -1,14 +1,13 @@
 #!/bin/bash
 #http://redsymbol.net/articles/unofficial-bash-strict-mode/
 
+source ~/.bashrc
+conda activate $BUILDKITE_BUILD_ID
+
 set -euo pipefail
 IFS=$'\n\t'
 
 set +e
-
-source ~/.bashrc
-conda activate $BUILDKITE_BUILD_ID
-
 set -x
 
 error=0

environments/linux-cuda/env.yml

@@ -1,10 +1,11 @@
 name: tmol
 channels:
-  - nvidia/label/cuda-12.1.1
+  - nvidia/label/cuda-12.8.0
   - conda-forge
 dependencies:
   - python=3.11
   - cuda
+  - nvtx
   - pip
   - pip:
     - -r requirements-linux-cuda.txt

environments/macos-cpu/env.yml

@@ -0,0 +1,8 @@
+name: tmol
+channels:
+  - conda-forge
+dependencies:
+  - python=3.11
+  - pip
+  - pip:
+    - -r requirements-macos-cpu.txt

environments/macos-cpu/requirements-dev-macos-cpu.txt

@@ -0,0 +1,181 @@
+#
+# This file is autogenerated by pip-compile with Python 3.11
+# by the following command:
+#
+#    pip-compile --constraint=requirements-macos-cpu.txt --no-emit-index-url --output-file=requirements-dev-macos-cpu.txt ../../requirements-dev.in
+#
+black==24.1.1
+    # via -r ../../requirements-dev.in
+build==1.0.3
+    # via pip-tools
+certifi==2024.2.2
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   requests
+cfgv==3.4.0
+    # via pre-commit
+charset-normalizer==3.3.2
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   requests
+clang-format==17.0.6
+    # via -r ../../requirements-dev.in
+click==8.1.7
+    # via
+    #   black
+    #   pip-tools
+codecov==2.1.13
+    # via -r ../../requirements-dev.in
+colorama==0.4.6
+    # via pytest-watch
+contourpy==1.2.0
+    # via matplotlib
+coverage[toml]==7.4.1
+    # via
+    #   codecov
+    #   pytest-cov
+cycler==0.12.1
+    # via matplotlib
+distlib==0.3.8
+    # via virtualenv
+docopt==0.6.2
+    # via pytest-watch
+filelock==3.13.1
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   virtualenv
+flake8==7.0.0
+    # via -r ../../requirements-dev.in
+fonttools==4.48.1
+    # via matplotlib
+identify==2.5.35
+    # via pre-commit
+idna==3.6
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   requests
+iniconfig==2.0.0
+    # via pytest
+itermplot==0.5
+    # via -r ../../requirements-dev.in
+kiwisolver==1.4.5
+    # via matplotlib
+matplotlib==3.8.2
+    # via
+    #   itermplot
+    #   seaborn
+mccabe==0.7.0
+    # via flake8
+mypy-extensions==1.0.0
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   black
+nodeenv==1.8.0
+    # via pre-commit
+numpy==1.26.4
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   contourpy
+    #   itermplot
+    #   matplotlib
+    #   pandas
+    #   seaborn
+packaging==23.2
+    # via
+    #   black
+    #   build
+    #   matplotlib
+    #   pytest
+pandas==2.2.0
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   seaborn
+pathspec==0.12.1
+    # via black
+pillow==10.2.0
+    # via matplotlib
+pip-tools==7.3.0
+    # via -r ../../requirements-dev.in
+platformdirs==4.2.0
+    # via
+    #   black
+    #   virtualenv
+pluggy==1.4.0
+    # via pytest
+pre-commit==3.6.2
+    # via -r ../../requirements-dev.in
+py==1.11.0
+    # via pytest-forked
+py-cpuinfo==9.0.0
+    # via pytest-benchmark
+pycodestyle==2.11.1
+    # via flake8
+pyflakes==3.2.0
+    # via flake8
+pyparsing==3.1.1
+    # via matplotlib
+pyproject-hooks==1.0.0
+    # via build
+pytest==8.0.0
+    # via
+    #   -r ../../requirements-dev.in
+    #   pytest-benchmark
+    #   pytest-cov
+    #   pytest-forked
+    #   pytest-instafail
+    #   pytest-repeat
+    #   pytest-watch
+pytest-benchmark==4.0.0
+    # via -r ../../requirements-dev.in
+pytest-cov==4.1.0
+    # via -r ../../requirements-dev.in
+pytest-forked==1.6.0
+    # via -r ../../requirements-dev.in
+pytest-instafail==0.5.0
+    # via -r ../../requirements-dev.in
+pytest-repeat==0.9.3
+    # via -r ../../requirements-dev.in
+pytest-watch==4.2.0
+    # via -r ../../requirements-dev.in
+python-dateutil==2.8.2
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   matplotlib
+    #   pandas
+pytz==2024.1
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   pandas
+pyyaml==6.0.1
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   pre-commit
+requests==2.31.0
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   codecov
+seaborn==0.13.2
+    # via -r ../../requirements-dev.in
+six==1.16.0
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   itermplot
+    #   python-dateutil
+tzdata==2023.4
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   pandas
+urllib3==2.2.0
+    # via
+    #   -c requirements-macos-cpu.txt
+    #   requests
+virtualenv==20.25.1
+    # via pre-commit
+watchdog==4.0.0
+    # via pytest-watch
+wheel==0.42.0
+    # via pip-tools
+
+# The following packages are considered to be unsafe in a requirements file:
+# pip
+# setuptools

environments/macos-cpu/requirements-macos-cpu.txt

@@ -0,0 +1,148 @@
+#
+# This file is autogenerated by pip-compile with Python 3.11
+# by the following command:
+#
+#    pip-compile --no-emit-index-url --output-file=requirements-macos-cpu.txt ../../requirements.in
+#
+asciitree==0.3.3
+    # via zarr
+astor==0.8.1
+    # via -r ../../requirements.in
+asttokens==2.4.1
+    # via stack-data
+attrs==23.2.0
+    # via
+    #   -r ../../requirements.in
+    #   attrs-strict
+    #   cattrs
+    #   hypothesis
+attrs-strict==1.0.1
+    # via -r ../../requirements.in
+cattrs==23.2.3
+    # via -r ../../requirements.in
+certifi==2024.2.2
+    # via requests
+charset-normalizer==3.3.2
+    # via requests
+decorator==5.1.1
+    # via
+    #   -r ../../requirements.in
+    #   ipython
+executing==2.0.1
+    # via stack-data
+fasteners==0.19
+    # via zarr
+filelock==3.13.1
+    # via torch
+frozendict==2.4.0
+    # via -r ../../requirements.in
+fsspec==2024.2.0
+    # via torch
+hypothesis==6.98.3
+    # via -r ../../requirements.in
+idna==3.6
+    # via requests
+ipython==8.21.0
+    # via -r ../../requirements.in
+jedi==0.19.1
+    # via ipython
+jinja2==3.1.3
+    # via torch
+llvmlite==0.42.0
+    # via numba
+markupsafe==2.1.5
+    # via jinja2
+matplotlib-inline==0.1.6
+    # via ipython
+mpmath==1.3.0
+    # via sympy
+mypy-extensions==1.0.0
+    # via typing-inspect
+networkx==3.2.1
+    # via
+    #   -r ../../requirements.in
+    #   torch
+ninja==1.11.1.1
+    # via -r ../../requirements.in
+numba==0.59.0
+    # via sparse
+numcodecs==0.12.1
+    # via zarr
+numpy==1.26.4
+    # via
+    #   -r ../../requirements.in
+    #   numba
+    #   numcodecs
+    #   pandas
+    #   pyarrow
+    #   scipy
+    #   sparse
+    #   zarr
+pandas==2.2.0
+    # via -r ../../requirements.in
+parso==0.8.3
+    # via jedi
+pexpect==4.9.0
+    # via ipython
+pint==0.23
+    # via -r ../../requirements.in
+prompt-toolkit==3.0.43
+    # via ipython
+psutil==5.9.8
+    # via -r ../../requirements.in
+ptyprocess==0.7.0
+    # via pexpect
+pure-eval==0.2.2
+    # via stack-data
+pyarrow==15.0.0
+    # via -r ../../requirements.in
+pygments==2.17.2
+    # via ipython
+python-dateutil==2.8.2
+    # via pandas
+pytz==2024.1
+    # via pandas
+pyyaml==6.0.1
+    # via -r ../../requirements.in
+requests==2.31.0
+    # via -r ../../requirements.in
+scipy==1.12.0
+    # via
+    #   -r ../../requirements.in
+    #   sparse
+six==1.16.0
+    # via
+    #   asttokens
+    #   python-dateutil
+sortedcontainers==2.4.0
+    # via hypothesis
+sparse==0.15.1
+    # via -r ../../requirements.in
+stack-data==0.6.3
+    # via ipython
+sympy==1.12
+    # via torch
+toolz==0.12.1
+    # via -r ../../requirements.in
+torch==2.2.0
+    # via -r ../../requirements.in
+traitlets==5.14.1
+    # via
+    #   ipython
+    #   matplotlib-inline
+typing-extensions==4.9.0
+    # via
+    #   -r ../../requirements.in
+    #   pint
+    #   torch
+    #   typing-inspect
+typing-inspect==0.9.0
+    # via -r ../../requirements.in
+tzdata==2023.4
+    # via pandas
+urllib3==2.2.0
+    # via requests
+wcwidth==0.2.13
+    # via prompt-toolkit
+zarr==2.16.1
+    # via -r ../../requirements.in

requirements.in

@@ -1,4 +1,4 @@
-torch == 2.6
+torch >= 2.5
 numpy # NOTE we get numpy from conda to ensure compat with numba
 astor
 attrs != 22.2.0 # problem unpickling tmol/tests/data/rosetta_baseline/1ubq.scores.pickle in 22.2.0 (see https://github.com/python-attrs/attrs/pull/1085)
@@ -18,6 +18,7 @@ requests
 scipy
 sparse # NOTE sparse>0.3.1 requires env var SPARSE_AUTO_DENSIFY=1 to be set
 toolz
+torchshow
 typing_extensions
 typing_inspect
 ninja