pytensor-ml 0.0.3__tar.gz → 0.0.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (93) hide show
  1. pytensor_ml-0.0.5/.gitattributes +1 -0
  2. pytensor_ml-0.0.5/.github/workflows/mypy.yml +57 -0
  3. pytensor_ml-0.0.5/.github/workflows/pre-commit.yml +16 -0
  4. pytensor_ml-0.0.5/.github/workflows/release.yml +50 -0
  5. pytensor_ml-0.0.5/.github/workflows/run_tests.yml +86 -0
  6. pytensor_ml-0.0.5/.gitignore +59 -0
  7. pytensor_ml-0.0.5/.pre-commit-config.yaml +65 -0
  8. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/PKG-INFO +26 -10
  9. pytensor_ml-0.0.5/conda_envs/pytensor_ml-gpu_jax.yml +42 -0
  10. pytensor_ml-0.0.5/conda_envs/pytensor_ml.yml +36 -0
  11. pytensor_ml-0.0.5/examples/mnist_feed_forward.ipynb +198 -0
  12. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/pyproject.toml +44 -18
  13. pytensor_ml-0.0.5/pytensor_ml/__init__.py +29 -0
  14. pytensor_ml-0.0.5/pytensor_ml/_version.py +24 -0
  15. pytensor_ml-0.0.5/pytensor_ml/activations.py +196 -0
  16. pytensor_ml-0.0.5/pytensor_ml/attention.py +253 -0
  17. pytensor_ml-0.0.5/pytensor_ml/checkpoint.py +127 -0
  18. pytensor_ml-0.0.5/pytensor_ml/dispatch/__init__.py +55 -0
  19. pytensor_ml-0.0.5/pytensor_ml/dispatch/jax/__init__.py +3 -0
  20. pytensor_ml-0.0.5/pytensor_ml/dispatch/jax/attention.py +25 -0
  21. pytensor_ml-0.0.5/pytensor_ml/dispatch/mlx/__init__.py +3 -0
  22. pytensor_ml-0.0.5/pytensor_ml/dispatch/mlx/attention.py +30 -0
  23. pytensor_ml-0.0.5/pytensor_ml/json_serialize.py +260 -0
  24. pytensor_ml-0.0.5/pytensor_ml/layers.py +474 -0
  25. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/pytensor_ml/loss.py +34 -0
  26. pytensor_ml-0.0.5/pytensor_ml/model.py +91 -0
  27. pytensor_ml-0.0.5/pytensor_ml/optim/__init__.py +37 -0
  28. pytensor_ml-0.0.5/pytensor_ml/optim/alias.py +119 -0
  29. pytensor_ml-0.0.5/pytensor_ml/optim/base.py +112 -0
  30. pytensor_ml-0.0.5/pytensor_ml/optim/clipping.py +64 -0
  31. pytensor_ml-0.0.5/pytensor_ml/optim/rules.py +279 -0
  32. pytensor_ml-0.0.5/pytensor_ml/optim/train.py +60 -0
  33. pytensor_ml-0.0.5/pytensor_ml/optim/transform.py +138 -0
  34. pytensor_ml-0.0.5/pytensor_ml/params.py +227 -0
  35. pytensor_ml-0.0.5/pytensor_ml/pretrained.py +254 -0
  36. pytensor_ml-0.0.5/pytensor_ml/pytensorf.py +326 -0
  37. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/pytensor_ml/rewriting/layers.py +1 -8
  38. pytensor_ml-0.0.5/pytensor_ml/serialize/__init__.py +2 -0
  39. pytensor_ml-0.0.5/pytensor_ml/serialize/blockwise.py +16 -0
  40. pytensor_ml-0.0.5/pytensor_ml/serialize/builders.py +45 -0
  41. pytensor_ml-0.0.5/pytensor_ml/serialize/elemwise.py +36 -0
  42. pytensor_ml-0.0.5/pytensor_ml/serialize/scalar.py +49 -0
  43. pytensor_ml-0.0.5/pytensor_ml/serialize/scan.py +42 -0
  44. pytensor_ml-0.0.5/pytensor_ml/state.py +131 -0
  45. pytensor_ml-0.0.5/pytensor_ml/transformer.py +170 -0
  46. pytensor_ml-0.0.5/scripts/run_mypy.py +12 -0
  47. pytensor_ml-0.0.5/tests/__init__.py +0 -0
  48. pytensor_ml-0.0.5/tests/conftest.py +43 -0
  49. pytensor_ml-0.0.5/tests/dispatch/__init__.py +0 -0
  50. pytensor_ml-0.0.5/tests/dispatch/jax/__init__.py +0 -0
  51. pytensor_ml-0.0.5/tests/dispatch/jax/test_attention.py +64 -0
  52. pytensor_ml-0.0.5/tests/dispatch/jax/test_basic.py +75 -0
  53. pytensor_ml-0.0.5/tests/dispatch/mlx/__init__.py +0 -0
  54. pytensor_ml-0.0.5/tests/dispatch/mlx/test_attention.py +65 -0
  55. pytensor_ml-0.0.5/tests/dispatch/mlx/test_basic.py +75 -0
  56. pytensor_ml-0.0.5/tests/optim/__init__.py +0 -0
  57. pytensor_ml-0.0.5/tests/optim/conftest.py +41 -0
  58. pytensor_ml-0.0.5/tests/optim/test_clipping.py +43 -0
  59. pytensor_ml-0.0.5/tests/optim/test_rules.py +118 -0
  60. pytensor_ml-0.0.5/tests/optim/test_training.py +138 -0
  61. pytensor_ml-0.0.5/tests/optim/test_transform.py +86 -0
  62. pytensor_ml-0.0.5/tests/rewriting/__init__.py +0 -0
  63. pytensor_ml-0.0.5/tests/rewriting/test_layers.py +62 -0
  64. pytensor_ml-0.0.5/tests/test_activations.py +85 -0
  65. pytensor_ml-0.0.5/tests/test_attention.py +243 -0
  66. pytensor_ml-0.0.5/tests/test_checkpoint.py +152 -0
  67. pytensor_ml-0.0.5/tests/test_layers.py +283 -0
  68. pytensor_ml-0.0.5/tests/test_loss.py +61 -0
  69. pytensor_ml-0.0.5/tests/test_model.py +116 -0
  70. pytensor_ml-0.0.5/tests/test_params.py +107 -0
  71. pytensor_ml-0.0.5/tests/test_pretrained.py +131 -0
  72. pytensor_ml-0.0.5/tests/test_pytensorf.py +38 -0
  73. pytensor_ml-0.0.5/tests/test_serialize.py +191 -0
  74. pytensor_ml-0.0.5/tests/test_state.py +34 -0
  75. pytensor_ml-0.0.5/tests/test_transformer.py +142 -0
  76. pytensor_ml-0.0.3/pytensor_ml/__init__.py +0 -18
  77. pytensor_ml-0.0.3/pytensor_ml/_version.py +0 -21
  78. pytensor_ml-0.0.3/pytensor_ml/activations.py +0 -57
  79. pytensor_ml-0.0.3/pytensor_ml/layers.py +0 -236
  80. pytensor_ml-0.0.3/pytensor_ml/model.py +0 -161
  81. pytensor_ml-0.0.3/pytensor_ml/optimizers.py +0 -387
  82. pytensor_ml-0.0.3/pytensor_ml/pytensorf.py +0 -115
  83. pytensor_ml-0.0.3/pytensor_ml.egg-info/PKG-INFO +0 -223
  84. pytensor_ml-0.0.3/pytensor_ml.egg-info/SOURCES.txt +0 -20
  85. pytensor_ml-0.0.3/pytensor_ml.egg-info/dependency_links.txt +0 -1
  86. pytensor_ml-0.0.3/pytensor_ml.egg-info/requires.txt +0 -10
  87. pytensor_ml-0.0.3/pytensor_ml.egg-info/top_level.txt +0 -1
  88. pytensor_ml-0.0.3/setup.cfg +0 -4
  89. pytensor_ml-0.0.3/setup.py +0 -19
  90. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/LICENSE +0 -0
  91. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/README.md +0 -0
  92. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/pytensor_ml/rewriting/__init__.py +0 -0
  93. {pytensor_ml-0.0.3 → pytensor_ml-0.0.5}/pytensor_ml/util.py +0 -0
@@ -0,0 +1 @@
1
+ pytensor_ml/_version.py export-subst
@@ -0,0 +1,57 @@
1
+ name: mypy
2
+
3
+ on:
4
+ pull_request:
5
+ push:
6
+ branches: [main]
7
+ paths:
8
+ - ".github/workflows/mypy.yml"
9
+ - "pytensor_ml/**"
10
+ - "scripts/run_mypy.py"
11
+ - "pyproject.toml"
12
+ - "conda_envs/**"
13
+
14
+ # Cancels all previous workflow runs for pull requests that have not completed.
15
+ concurrency:
16
+ group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
17
+ cancel-in-progress: true
18
+
19
+ jobs:
20
+ mypy:
21
+ runs-on: ubuntu-latest
22
+
23
+ defaults:
24
+ run:
25
+ shell: bash -l {0}
26
+
27
+ steps:
28
+ - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
29
+ - uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
30
+ env:
31
+ CACHE_NUMBER: 0
32
+ with:
33
+ path: ~/conda_pkgs_dir
34
+ key: ${{ runner.os }}-py312-conda-${{ env.CACHE_NUMBER }}-${{
35
+ hashFiles('conda_envs/pytensor_ml.yml') }}
36
+ - uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1
37
+ with:
38
+ miniforge-variant: Miniforge3
39
+ miniforge-version: latest
40
+ mamba-version: "*"
41
+ activate-environment: pytensor_ml
42
+ channel-priority: strict
43
+ environment-file: conda_envs/pytensor_ml.yml
44
+ python-version: 3.12
45
+ use-mamba: true
46
+ use-only-tar-bz2: false
47
+
48
+ - name: Install current branch
49
+ run: |
50
+ conda activate pytensor_ml
51
+ pip install -e .
52
+ python --version
53
+
54
+ - name: Run mypy
55
+ run: |
56
+ conda activate pytensor_ml
57
+ python ./scripts/run_mypy.py
@@ -0,0 +1,16 @@
1
+ name: pre-commit
2
+
3
+ on:
4
+ pull_request:
5
+ push:
6
+ branches: [main]
7
+
8
+ jobs:
9
+ pre-commit:
10
+ runs-on: ubuntu-latest
11
+ env:
12
+ SKIP: no-commit-to-branch
13
+ steps:
14
+ - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
15
+ - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
16
+ - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
@@ -0,0 +1,50 @@
1
+ name: release-pipeline
2
+
3
+ permissions: {}
4
+
5
+ on:
6
+ push:
7
+ branches:
8
+ - main
9
+ release:
10
+ types:
11
+ - published
12
+
13
+ jobs:
14
+ build-package:
15
+ runs-on: ubuntu-latest
16
+ permissions:
17
+ # write attestations and id-token are necessary for attest-build-provenance-github
18
+ attestations: write
19
+ id-token: write
20
+ steps:
21
+ - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
22
+ with:
23
+ # hatch-vcs derives the version from git tags, so the full history is required.
24
+ fetch-depth: 0
25
+ persist-credentials: false
26
+ - uses: hynek/build-and-inspect-python-package@d44ca7d91762de7a7d5436ddae667c6da6d1c3df # v2.18.0
27
+ with:
28
+ # Prove that the packages were built in the context of this workflow.
29
+ attest-build-provenance-github: true
30
+
31
+ publish-package:
32
+ # Don't publish from forks
33
+ if: github.repository_owner == 'pymc-devs' && github.event_name == 'release' && github.event.action == 'published'
34
+ # Use the `release` GitHub environment to protect the Trusted Publishing (OIDC)
35
+ # workflow by requiring signoff from a maintainer.
36
+ environment: release
37
+ needs: build-package
38
+ runs-on: ubuntu-latest
39
+ permissions:
40
+ # write id-token is necessary for trusted publishing (OIDC)
41
+ id-token: write
42
+ steps:
43
+ - name: Download Distribution Artifacts
44
+ uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
45
+ with:
46
+ # The build-and-inspect-python-package action invokes upload-artifact.
47
+ name: Packages
48
+ path: dist
49
+ - name: Publish Package to PyPI
50
+ uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0
@@ -0,0 +1,86 @@
1
+ name: run_tests
2
+
3
+ on:
4
+ pull_request:
5
+ push:
6
+ branches: [main]
7
+
8
+ # Cancels all previous workflow runs for pull requests that have not completed.
9
+ concurrency:
10
+ # The concurrency group contains the workflow name and the branch name for pull requests
11
+ # or the commit hash for any other events.
12
+ group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
13
+ cancel-in-progress: true
14
+
15
+ jobs:
16
+ unittest:
17
+ strategy:
18
+ fail-fast: false
19
+ matrix:
20
+ os: [ubuntu-latest, windows-latest]
21
+ python-version: ["3.12"]
22
+ test-subset:
23
+ - |
24
+ tests/
25
+
26
+ runs-on: ${{ matrix.os }}
27
+
28
+ env:
29
+ TEST_SUBSET: ${{ matrix.test-subset }}
30
+
31
+ defaults:
32
+ run:
33
+ shell: bash -l {0}
34
+
35
+ steps:
36
+ - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
37
+ - uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
38
+ env:
39
+ CACHE_NUMBER: 0
40
+ with:
41
+ path: ~/conda_pkgs_dir
42
+ key: ${{ runner.os }}-py${{matrix.python-version}}-conda-${{ env.CACHE_NUMBER }}-${{
43
+ hashFiles('conda_envs/pytensor_ml.yml') }}
44
+ - name: Cache multiple paths
45
+ uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
46
+ env:
47
+ # Increase this value to reset cache if pyproject.toml has not changed
48
+ CACHE_NUMBER: 0
49
+ with:
50
+ path: |
51
+ ~/.cache/pip
52
+ $RUNNER_TOOL_CACHE/Python/*
53
+ ~\AppData\Local\pip\Cache
54
+ key: ${{ runner.os }}-build-${{ matrix.python-version }}-${{ env.CACHE_NUMBER }}-${{
55
+ hashFiles('pyproject.toml') }}
56
+ - uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1
57
+ with:
58
+ miniforge-variant: Miniforge3
59
+ miniforge-version: latest
60
+ mamba-version: "*"
61
+ activate-environment: pytensor_ml
62
+ channel-priority: strict
63
+ environment-file: conda_envs/pytensor_ml.yml
64
+ python-version: 3.12
65
+ use-mamba: true
66
+
67
+ # IMPORTANT: This may break caching of conda packages! See https://github.com/conda-incubator/setup-miniconda/issues/267
68
+ use-only-tar-bz2: false
69
+
70
+ - name: Install current branch
71
+ run: |
72
+ conda activate pytensor_ml
73
+ pip install -e .
74
+ python --version
75
+
76
+ - name: Run tests
77
+ run: |
78
+ python -m pytest -vv --cov=pytensor_ml --cov-report=xml --no-cov-on-fail --cov-report term $TEST_SUBSET
79
+ - name: Upload coverage to Codecov
80
+ uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1
81
+ with:
82
+ token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
83
+ env_vars: TEST_SUBSET
84
+ name: ${{ matrix.os }}
85
+ fail_ci_if_error: false
86
+ verbose: true
@@ -0,0 +1,59 @@
1
+ *.pkl
2
+ _build
3
+ __pycache__
4
+ .coverage
5
+ *.linkinfo
6
+ *.o
7
+ *.c
8
+ *.orig
9
+ *.pyc
10
+ *.pyo
11
+ *.so
12
+ *.sw?
13
+ *~
14
+ *.aux
15
+ *.log
16
+ *.nav
17
+ *.out
18
+ *.snm
19
+ *.toc
20
+ *.vrb
21
+ *.nbc
22
+ *.nbi
23
+ .noseids
24
+ *.DS_Store
25
+ *.bak
26
+ *.egg-info/
27
+ \#*\#
28
+ build
29
+ compiled/*.cpp
30
+ core.*
31
+ cutils_ext.cpp
32
+ dist
33
+ doc/.build/
34
+ doc/indexes/oplist.txt
35
+ doc/indexes/typelist.txt
36
+ html
37
+ pdf
38
+ setuptools-*.egg
39
+ pytensor/generated_version.py
40
+ pytensor/generated_version.py.out
41
+ distribute-*.egg
42
+ distribute-*.tar.gz
43
+ PyTensor.suo
44
+ .ipynb_checkpoints
45
+ .pydevproject
46
+ .ropeproject
47
+ core
48
+ .idea
49
+ .vs
50
+ .mypy_cache/
51
+ /htmlcov/
52
+ pytensor_ml/_version.py
53
+
54
+ pytensor-venv/
55
+ /notebooks/Sandbox*
56
+ .vscode/
57
+ testing-report.html
58
+ coverage.xml
59
+ .coverage.*
@@ -0,0 +1,65 @@
1
+ repos:
2
+ - repo: https://github.com/pre-commit/pre-commit-hooks
3
+ rev: v6.0.0
4
+ hooks:
5
+ - id: check-merge-conflict
6
+ - id: check-toml
7
+ - id: check-yaml
8
+ - id: debug-statements
9
+ - id: end-of-file-fixer
10
+ - id: no-commit-to-branch
11
+ args: [--branch, main]
12
+ - id: trailing-whitespace
13
+
14
+ - repo: https://github.com/astral-sh/ruff-pre-commit
15
+ rev: v0.15.17
16
+ hooks:
17
+ - id: ruff
18
+ args: [ --fix, --exit-non-zero-on-fix ]
19
+ - id: ruff-format
20
+ types_or: [ python, pyi, jupyter ]
21
+
22
+ - repo: https://github.com/MarcoGorelli/madforhooks
23
+ rev: 0.4.1
24
+ hooks:
25
+ - id: no-print-statements
26
+ types: [python]
27
+ exclude: |
28
+ (?x)^
29
+ |pytensor_ml/_version.py
30
+ - id: check-execution-order
31
+ args: [--strict]
32
+
33
+ - repo: local
34
+ hooks:
35
+ - id: notebook_name
36
+ entry: '\(notebook_name\)='
37
+ language: pygrep
38
+ minimum_pre_commit_version: 2.8.0
39
+ name: Check notebooks do not use literally notebook_name as target
40
+ types: [jupyter]
41
+
42
+ - id: no-references-as-links
43
+ name: Check no references that should be sphinx cross-references are urls
44
+ description: >-
45
+ 'A quick check to prevent urls pointing to pymc docs or other sphinx built docs like arviz, numpy, scipy...'
46
+ files: ^examples/.+\.ipynb$
47
+ exclude: >
48
+ (?x)(index.md|
49
+ 404.md|
50
+ conf.py)
51
+ entry: >
52
+ (?x)(arviz-devs.github.io|
53
+ pymc-extras.readthedocs.io|
54
+ docs.pymc.io|
55
+ numpy.org/doc|
56
+ pymc-examples.readthedocs.io|
57
+ docs.python.org|
58
+ xarray.pydata.org|
59
+ python.arviz.org|
60
+ pytensor.readthedocs.io|
61
+ docs.xarray.dev|
62
+ www.pymc.io|
63
+ docs.scipy.org/doc)
64
+ language: pygrep
65
+ types_or: [markdown, rst, jupyter]
@@ -1,7 +1,10 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: pytensor_ml
3
- Version: 0.0.3
3
+ Version: 0.0.5
4
4
  Summary: A(nother) deep learning library, built on top of Pytensor
5
+ Project-URL: Homepage, https://github.com/pymc-devs/pytensor_ml
6
+ Project-URL: Source, https://github.com/pymc-devs/pytensor_ml
7
+ Project-URL: Tracker, https://github.com/pymc-devs/pytensor_ml/issues
5
8
  Author: Ricardo Vieira
6
9
  Author-email: Jesse Grabowski <jessegrabowski@gmail.com>
7
10
  License: Apache License
@@ -205,19 +208,32 @@ License: Apache License
205
208
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
206
209
  See the License for the specific language governing permissions and
207
210
  limitations under the License.
208
-
209
- Keywords: neural network,deep learning,autograd,gradient descent,pymc
211
+ License-File: LICENSE
212
+ Keywords: autograd,deep learning,gradient descent,neural network,pymc
210
213
  Classifier: Development Status :: 2 - Pre-Alpha
211
214
  Classifier: Intended Audience :: Science/Research
212
- Classifier: Programming Language :: Python
213
- Classifier: Topic :: Scientific/Engineering
215
+ Classifier: Operating System :: MacOS
214
216
  Classifier: Operating System :: Microsoft :: Windows
215
217
  Classifier: Operating System :: POSIX
216
218
  Classifier: Operating System :: Unix
217
- Classifier: Operating System :: MacOS
219
+ Classifier: Programming Language :: Python
218
220
  Classifier: Programming Language :: Python :: 3
219
- Classifier: Programming Language :: Python :: 3.10
220
- Classifier: Programming Language :: Python :: 3.11
221
221
  Classifier: Programming Language :: Python :: 3.12
222
+ Classifier: Programming Language :: Python :: 3.13
223
+ Classifier: Programming Language :: Python :: 3.14
224
+ Classifier: Topic :: Scientific/Engineering
225
+ Requires-Python: >=3.12
226
+ Requires-Dist: numpy
227
+ Requires-Dist: pytensor<4.0.0,>=3.2.3
228
+ Requires-Dist: safetensors
222
229
  Provides-Extra: dev
223
- License-File: LICENSE
230
+ Requires-Dist: jax; extra == 'dev'
231
+ Requires-Dist: mypy; extra == 'dev'
232
+ Requires-Dist: pre-commit; extra == 'dev'
233
+ Requires-Dist: pytest; extra == 'dev'
234
+ Requires-Dist: pytest-cov; extra == 'dev'
235
+ Requires-Dist: scikit-learn; extra == 'dev'
236
+ Description-Content-Type: text/markdown
237
+
238
+ # pytensor_ml
239
+ Neural network package built on Pytensor
@@ -0,0 +1,42 @@
1
+ name: pytensor_ml
2
+ channels:
3
+ - conda-forge
4
+ - nvidia
5
+ - nodefaults
6
+
7
+ dependencies:
8
+ - python>=3.12
9
+ - pytensor>=3.2.3,<4.0.0
10
+ - numpy
11
+ - scikit-learn
12
+
13
+ # GPU Stuff
14
+ - jax
15
+ - numpyro
16
+ - blackjax
17
+ - jaxlib=*=*cuda*
18
+ - cuda-nvcc
19
+
20
+ # For testing
21
+ - coveralls
22
+ - diff-cover
23
+ - pytest
24
+ - pytest-cov
25
+ - pytest-xdist
26
+ - pytest-benchmark
27
+ - pytest-mock
28
+
29
+ # For building docs
30
+ - sphinx
31
+ - sphinx_rtd_theme
32
+ - pygments
33
+ - pydot
34
+ - ipython
35
+
36
+ # code style
37
+ - ruff
38
+
39
+ # developer tools
40
+ - pre-commit
41
+ - packaging
42
+ - typing_extensions
@@ -0,0 +1,36 @@
1
+ name: pytensor_ml
2
+ channels:
3
+ - conda-forge
4
+ - nodefaults
5
+
6
+ dependencies:
7
+ - python>=3.12
8
+ - pytensor>=3.2.3,<4.0.0
9
+ - numpy
10
+ - safetensors
11
+ - scikit-learn
12
+
13
+ # For testing
14
+ - coveralls
15
+ - diff-cover
16
+ - pytest
17
+ - pytest-cov
18
+ - pytest-xdist
19
+ - pytest-benchmark
20
+ - pytest-mock
21
+
22
+ # For building docs
23
+ - sphinx
24
+ - sphinx_rtd_theme
25
+ - pygments
26
+ - pydot
27
+ - ipython
28
+
29
+ # developer tools
30
+ - pre-commit
31
+ - mypy
32
+
33
+ # JAX backend for dispatch tests; via pip so Windows pulls CPU wheels instead of conda jaxlib
34
+ - pip
35
+ - pip:
36
+ - jax
@@ -0,0 +1,198 @@
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "code",
5
+ "execution_count": null,
6
+ "id": "15279e7f",
7
+ "metadata": {},
8
+ "outputs": [],
9
+ "source": [
10
+ "%load_ext autoreload\n",
11
+ "%autoreload 2"
12
+ ]
13
+ },
14
+ {
15
+ "cell_type": "code",
16
+ "execution_count": null,
17
+ "id": "00915a61",
18
+ "metadata": {},
19
+ "outputs": [],
20
+ "source": [
21
+ "import pytensor\n",
22
+ "\n",
23
+ "floatX = \"float32\"\n",
24
+ "pytensor.config.floatX = floatX"
25
+ ]
26
+ },
27
+ {
28
+ "cell_type": "code",
29
+ "execution_count": null,
30
+ "id": "fae24c05",
31
+ "metadata": {},
32
+ "outputs": [],
33
+ "source": [
34
+ "import numpy as np\n",
35
+ "\n",
36
+ "from sklearn.datasets import load_digits\n",
37
+ "from sklearn.preprocessing import MinMaxScaler, OneHotEncoder\n",
38
+ "\n",
39
+ "from pytensor_ml.activations import LeakyReLU\n",
40
+ "from pytensor_ml.layers import BatchNorm2D, Dropout, Linear, Sequential\n",
41
+ "from pytensor_ml.loss import CrossEntropy\n",
42
+ "from pytensor_ml.model import Model\n",
43
+ "from pytensor_ml.optim import adagrad\n",
44
+ "from pytensor_ml.util import DataLoader"
45
+ ]
46
+ },
47
+ {
48
+ "cell_type": "code",
49
+ "execution_count": null,
50
+ "id": "65882b64",
51
+ "metadata": {},
52
+ "outputs": [],
53
+ "source": [
54
+ "X, y = load_digits(return_X_y=True)\n",
55
+ "y_onehot = OneHotEncoder().fit_transform(y[:, None]).toarray().astype(floatX)\n",
56
+ "X_normed = MinMaxScaler().fit_transform(X).astype(floatX)"
57
+ ]
58
+ },
59
+ {
60
+ "cell_type": "code",
61
+ "execution_count": null,
62
+ "id": "f81f9076",
63
+ "metadata": {},
64
+ "outputs": [],
65
+ "source": [
66
+ "# X_in = Input('X_in', shape=(64, ))\n",
67
+ "X_in = pytensor.tensor.tensor(\"X_in\", shape=(None, 64))"
68
+ ]
69
+ },
70
+ {
71
+ "cell_type": "code",
72
+ "execution_count": null,
73
+ "id": "41840fac",
74
+ "metadata": {},
75
+ "outputs": [],
76
+ "source": [
77
+ "prediction_network = Sequential(\n",
78
+ " Linear(\"Linear_1\", n_in=64, n_out=256),\n",
79
+ " BatchNorm2D(\"BatchNorm_1\"),\n",
80
+ " Dropout(\"Dropout_1\", p=0.5),\n",
81
+ " LeakyReLU(),\n",
82
+ " Linear(\"Linear_2\", n_in=256, n_out=128),\n",
83
+ " BatchNorm2D(\"BatchNorm_2\"),\n",
84
+ " Dropout(\"Dropout_2\", p=0.5),\n",
85
+ " LeakyReLU(),\n",
86
+ " Linear(\"Logits\", n_in=128, n_out=10),\n",
87
+ ")\n",
88
+ "\n",
89
+ "y_hat = prediction_network(X_in)\n",
90
+ "model = Model(X_in, y_hat, compile_kwargs={\"mode\": \"NUMBA\"})\n",
91
+ "model"
92
+ ]
93
+ },
94
+ {
95
+ "cell_type": "code",
96
+ "execution_count": null,
97
+ "id": "3bfa0ddc",
98
+ "metadata": {},
99
+ "outputs": [],
100
+ "source": [
101
+ "loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction=\"mean\")"
102
+ ]
103
+ },
104
+ {
105
+ "cell_type": "code",
106
+ "execution_count": null,
107
+ "id": "290d3191",
108
+ "metadata": {},
109
+ "outputs": [],
110
+ "source": [
111
+ "# `model.compile_train` is the training companion to `model.predict`: it builds the loss against a target,\n",
112
+ "# differentiates it, folds in the BatchNorm running-statistic updates, and compiles a one-step training\n",
113
+ "# function (reusing the model's NUMBA compile settings).\n",
114
+ "step = model.compile_train(adagrad(learning_rate=1e-3), loss_fn, ndim_out=2)"
115
+ ]
116
+ },
117
+ {
118
+ "cell_type": "code",
119
+ "execution_count": null,
120
+ "id": "e77b697f",
121
+ "metadata": {},
122
+ "outputs": [],
123
+ "source": [
124
+ "model.initialize()"
125
+ ]
126
+ },
127
+ {
128
+ "cell_type": "code",
129
+ "execution_count": null,
130
+ "id": "dd6dede5",
131
+ "metadata": {},
132
+ "outputs": [],
133
+ "source": [
134
+ "from tqdm.notebook import tqdm\n",
135
+ "\n",
136
+ "loader = DataLoader(X_normed, y_onehot)\n",
137
+ "loss_history = []\n",
138
+ "n_epochs = 10_000\n",
139
+ "\n",
140
+ "for _ in tqdm(range(n_epochs)):\n",
141
+ " X_batch, y_batch = loader()\n",
142
+ " loss_value = step(X_batch, y_batch)\n",
143
+ " loss_history.append(loss_value)"
144
+ ]
145
+ },
146
+ {
147
+ "cell_type": "code",
148
+ "execution_count": null,
149
+ "id": "ab1ad6ef",
150
+ "metadata": {},
151
+ "outputs": [],
152
+ "source": [
153
+ "import matplotlib.pyplot as plt\n",
154
+ "\n",
155
+ "plt.plot(loss_history)"
156
+ ]
157
+ },
158
+ {
159
+ "cell_type": "code",
160
+ "execution_count": null,
161
+ "id": "fb38dd99",
162
+ "metadata": {},
163
+ "outputs": [],
164
+ "source": [
165
+ "from scipy.special import softmax\n",
166
+ "\n",
167
+ "# `model.predict` compiles a forward pass specialized for inference: Dropout is dropped and BatchNorm uses its\n",
168
+ "# running statistics.\n",
169
+ "y_hat_logits = model.predict(X_normed)\n",
170
+ "y_hat_probs = softmax(y_hat_logits, axis=-1)\n",
171
+ "y_pred = np.argmax(y_hat_probs, axis=-1)"
172
+ ]
173
+ },
174
+ {
175
+ "cell_type": "code",
176
+ "execution_count": null,
177
+ "id": "4393b9c3",
178
+ "metadata": {},
179
+ "outputs": [],
180
+ "source": [
181
+ "import seaborn as sns\n",
182
+ "\n",
183
+ "from sklearn.metrics import confusion_matrix\n",
184
+ "\n",
185
+ "sns.heatmap(confusion_matrix(y, y_pred), annot=True, fmt=\"0.0f\")"
186
+ ]
187
+ }
188
+ ],
189
+ "metadata": {
190
+ "kernelspec": {
191
+ "display_name": "Python 3 (ipykernel)",
192
+ "language": "python",
193
+ "name": "python3"
194
+ }
195
+ },
196
+ "nbformat": 4,
197
+ "nbformat_minor": 5
198
+ }