pytensor-ml 0.2.1__tar.gz → 0.2.2__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 (171) hide show
  1. pytensor_ml-0.2.2/.github/workflows/rtd-preview.yml +30 -0
  2. pytensor_ml-0.2.2/.readthedocs.yaml +24 -0
  3. pytensor_ml-0.2.2/PKG-INFO +138 -0
  4. pytensor_ml-0.2.2/README.md +96 -0
  5. pytensor_ml-0.2.2/conda_envs/environment-docs.yml +35 -0
  6. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/conda_envs/pytensor_ml.yml +0 -7
  7. pytensor_ml-0.2.2/docs/.gitignore +16 -0
  8. pytensor_ml-0.2.2/docs/Makefile +34 -0
  9. pytensor_ml-0.2.2/docs/make.bat +29 -0
  10. pytensor_ml-0.2.2/docs/source/_templates/autosummary/class.rst +34 -0
  11. pytensor_ml-0.2.2/docs/source/api/activations.rst +17 -0
  12. pytensor_ml-0.2.2/docs/source/api/layers.rst +101 -0
  13. pytensor_ml-0.2.2/docs/source/api/loss.rst +12 -0
  14. pytensor_ml-0.2.2/docs/source/api/model.rst +9 -0
  15. pytensor_ml-0.2.2/docs/source/api/optim.rst +133 -0
  16. pytensor_ml-0.2.2/docs/source/api/pytensorf.rst +33 -0
  17. pytensor_ml-0.2.2/docs/source/api/serialization.rst +24 -0
  18. pytensor_ml-0.2.2/docs/source/api/state.rst +37 -0
  19. pytensor_ml-0.2.2/docs/source/api/util.rst +9 -0
  20. pytensor_ml-0.2.2/docs/source/api.rst +18 -0
  21. pytensor_ml-0.2.2/docs/source/conf.py +176 -0
  22. pytensor_ml-0.2.2/docs/source/dev/contributing.rst +91 -0
  23. pytensor_ml-0.2.2/docs/source/dev/docs.rst +87 -0
  24. pytensor_ml-0.2.2/docs/source/dev/index.rst +9 -0
  25. pytensor_ml-0.2.2/docs/source/dev/style_guide.rst +246 -0
  26. pytensor_ml-0.2.2/docs/source/index.rst +81 -0
  27. pytensor_ml-0.2.2/docs/source/install.rst +87 -0
  28. pytensor_ml-0.2.2/docs/source/references.bib +44 -0
  29. pytensor_ml-0.2.2/docs/source/sharp_edges.rst +247 -0
  30. pytensor_ml-0.2.2/docs/sphinxext/generate_gallery.py +318 -0
  31. pytensor_ml-0.2.2/examples/gallery/mnist_feed_forward.ipynb +758 -0
  32. pytensor_ml-0.2.2/pixi.lock +10749 -0
  33. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pyproject.toml +49 -0
  34. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/_version.py +2 -2
  35. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/activations.py +155 -1
  36. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/base.py +29 -2
  37. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/checkpoint.py +39 -0
  38. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/json_serialize.py +1 -1
  39. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/attention.py +44 -1
  40. pytensor_ml-0.2.2/pytensor_ml/layers/combinators.py +171 -0
  41. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/conv.py +124 -9
  42. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/dropout.py +19 -0
  43. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/embedding.py +12 -0
  44. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/linear.py +17 -0
  45. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/norm.py +37 -0
  46. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/padding.py +134 -1
  47. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/recurrent.py +143 -4
  48. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/transformer.py +31 -0
  49. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/loss.py +105 -0
  50. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/model.py +34 -3
  51. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/optim/__init__.py +14 -2
  52. pytensor_ml-0.2.2/pytensor_ml/optim/alias.py +507 -0
  53. pytensor_ml-0.2.2/pytensor_ml/optim/base.py +742 -0
  54. pytensor_ml-0.2.2/pytensor_ml/optim/checks.py +66 -0
  55. pytensor_ml-0.2.2/pytensor_ml/optim/clipping.py +153 -0
  56. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/optim/guards.py +155 -19
  57. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/optim/policy.py +89 -31
  58. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/optim/rules.py +264 -60
  59. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/optim/schedules.py +154 -49
  60. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/optim/train.py +41 -6
  61. pytensor_ml-0.2.2/pytensor_ml/optim/transform.py +284 -0
  62. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/params.py +93 -2
  63. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/pretrained.py +76 -0
  64. pytensor_ml-0.2.2/pytensor_ml/pytensorf/collect.py +409 -0
  65. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/pytensorf/compile.py +63 -4
  66. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/pytensorf/rewrite.py +36 -2
  67. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/pytensorf/rng.py +20 -3
  68. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/rewriting/scan.py +1 -1
  69. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/state.py +208 -3
  70. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/util.py +17 -0
  71. pytensor_ml-0.2.2/tests/optim/test_clipping.py +126 -0
  72. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_guards.py +53 -0
  73. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_policy.py +90 -0
  74. pytensor_ml-0.2.2/tests/optim/test_public_surface.py +25 -0
  75. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_schedules.py +142 -1
  76. pytensor_ml-0.2.2/tests/optim/test_update_space.py +479 -0
  77. pytensor_ml-0.2.2/tests/rewriting/__init__.py +0 -0
  78. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_collect.py +11 -0
  79. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_pytensorf.py +113 -1
  80. pytensor_ml-0.2.1/PKG-INFO +0 -43
  81. pytensor_ml-0.2.1/README.md +0 -2
  82. pytensor_ml-0.2.1/examples/mnist_feed_forward.ipynb +0 -197
  83. pytensor_ml-0.2.1/pytensor_ml/layers/combinators.py +0 -51
  84. pytensor_ml-0.2.1/pytensor_ml/optim/alias.py +0 -314
  85. pytensor_ml-0.2.1/pytensor_ml/optim/base.py +0 -305
  86. pytensor_ml-0.2.1/pytensor_ml/optim/clipping.py +0 -64
  87. pytensor_ml-0.2.1/pytensor_ml/optim/transform.py +0 -153
  88. pytensor_ml-0.2.1/pytensor_ml/pytensorf/collect.py +0 -192
  89. pytensor_ml-0.2.1/tests/optim/test_clipping.py +0 -43
  90. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/.gitattributes +0 -0
  91. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/.github/workflows/mypy.yml +0 -0
  92. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/.github/workflows/pre-commit.yml +0 -0
  93. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/.github/workflows/release.yml +0 -0
  94. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/.github/workflows/run_tests.yml +0 -0
  95. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/.gitignore +0 -0
  96. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/.pre-commit-config.yaml +0 -0
  97. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/LICENSE +0 -0
  98. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/conda_envs/pytensor_ml-gpu_jax.yml +0 -0
  99. /pytensor_ml-0.2.1/pytensor_ml/rewriting/__init__.py → /pytensor_ml-0.2.2/docs/source/_static/.gitkeep +0 -0
  100. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/__init__.py +0 -0
  101. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/__init__.py +0 -0
  102. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/jax/__init__.py +0 -0
  103. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/jax/attention.py +0 -0
  104. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/jax/conv.py +0 -0
  105. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/jax/pooling.py +0 -0
  106. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/mlx/__init__.py +0 -0
  107. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/mlx/attention.py +0 -0
  108. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/mlx/conv.py +0 -0
  109. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/mlx/pooling.py +0 -0
  110. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/numba/__init__.py +0 -0
  111. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/numba/conv.py +0 -0
  112. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/pytorch/__init__.py +0 -0
  113. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/pytorch/conv.py +0 -0
  114. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/dispatch/pytorch/pooling.py +0 -0
  115. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/layers/__init__.py +0 -0
  116. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/pytensorf/__init__.py +0 -0
  117. {pytensor_ml-0.2.1/tests → pytensor_ml-0.2.2/pytensor_ml/rewriting}/__init__.py +0 -0
  118. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/rewriting/conv.py +0 -0
  119. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/rewriting/layers.py +0 -0
  120. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/serialize/__init__.py +0 -0
  121. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/serialize/base.py +0 -0
  122. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/serialize/blockwise.py +0 -0
  123. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/serialize/builders.py +0 -0
  124. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/serialize/elemwise.py +0 -0
  125. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/serialize/scalar.py +0 -0
  126. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/pytensor_ml/serialize/scan.py +0 -0
  127. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/scripts/run_mypy.py +0 -0
  128. {pytensor_ml-0.2.1/tests/dispatch → pytensor_ml-0.2.2/tests}/__init__.py +0 -0
  129. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/conftest.py +0 -0
  130. {pytensor_ml-0.2.1/tests/dispatch/jax → pytensor_ml-0.2.2/tests/dispatch}/__init__.py +0 -0
  131. {pytensor_ml-0.2.1/tests/dispatch/mlx → pytensor_ml-0.2.2/tests/dispatch/jax}/__init__.py +0 -0
  132. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/jax/test_attention.py +0 -0
  133. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/jax/test_basic.py +0 -0
  134. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/jax/test_conv.py +0 -0
  135. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/jax/test_pool.py +0 -0
  136. {pytensor_ml-0.2.1/tests/dispatch/pytorch → pytensor_ml-0.2.2/tests/dispatch/mlx}/__init__.py +0 -0
  137. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/mlx/test_attention.py +0 -0
  138. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/mlx/test_basic.py +0 -0
  139. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/mlx/test_conv.py +0 -0
  140. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/mlx/test_pool.py +0 -0
  141. {pytensor_ml-0.2.1/tests/optim → pytensor_ml-0.2.2/tests/dispatch/pytorch}/__init__.py +0 -0
  142. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/pytorch/test_conv.py +0 -0
  143. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/dispatch/pytorch/test_pool.py +0 -0
  144. {pytensor_ml-0.2.1/tests/rewriting → pytensor_ml-0.2.2/tests/optim}/__init__.py +0 -0
  145. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/conftest.py +0 -0
  146. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_composition.py +0 -0
  147. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_learning_rate.py +0 -0
  148. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_rules.py +0 -0
  149. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_training.py +0 -0
  150. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/optim/test_transform.py +0 -0
  151. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/rewriting/conftest.py +0 -0
  152. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/rewriting/test_layers.py +0 -0
  153. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/rewriting/test_scan.py +0 -0
  154. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_activations.py +0 -0
  155. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_attention.py +0 -0
  156. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_base.py +0 -0
  157. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_checkpoint.py +0 -0
  158. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_conv.py +0 -0
  159. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_conv_transpose.py +0 -0
  160. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_layers.py +0 -0
  161. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_loss.py +0 -0
  162. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_model.py +0 -0
  163. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_padding.py +0 -0
  164. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_pool.py +0 -0
  165. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_pretrained.py +0 -0
  166. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_recurrent.py +0 -0
  167. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_serialize.py +0 -0
  168. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_state.py +0 -0
  169. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_transformer.py +0 -0
  170. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_util.py +0 -0
  171. {pytensor_ml-0.2.1 → pytensor_ml-0.2.2}/tests/test_workflow_groups.py +0 -0
@@ -0,0 +1,30 @@
1
+ name: Read the Docs PR preview
2
+
3
+ on:
4
+ pull_request_target:
5
+ types:
6
+ - opened
7
+ - synchronize
8
+ - reopened
9
+ # Only fire when something that could change the docs build is touched.
10
+ paths:
11
+ - "docs/**"
12
+ - ".readthedocs.yaml"
13
+ - "conda_envs/environment-docs.yml"
14
+ - "examples/**"
15
+ - "pytensor_ml/**"
16
+ - "pyproject.toml"
17
+
18
+ permissions:
19
+ pull-requests: write
20
+
21
+ jobs:
22
+ documentation-links:
23
+ runs-on: ubuntu-latest
24
+ steps:
25
+ - uses: readthedocs/actions/preview@v1
26
+ with:
27
+ # Project slug as configured on Read the Docs. Verify this matches the slug shown in the
28
+ # RTD project URL (https://readthedocs.org/projects/<slug>/) once the project is imported
29
+ # there.
30
+ project-slug: "pytensor-ml"
@@ -0,0 +1,24 @@
1
+ version: 2
2
+
3
+ sphinx:
4
+ configuration: docs/source/conf.py
5
+ fail_on_warning: false
6
+
7
+ conda:
8
+ environment: conda_envs/environment-docs.yml
9
+
10
+ python:
11
+ install:
12
+ - method: pip
13
+ path: .
14
+
15
+ build:
16
+ os: "ubuntu-22.04"
17
+ tools:
18
+ python: "miniforge3-latest"
19
+ jobs:
20
+ # hatch-vcs derives the version from git tags; RTD's default checkout is shallow and tagless, so
21
+ # fetch tags or pytensor_ml.__version__ falls back to 0.0.0+unknown and the version selector
22
+ # labels break.
23
+ post_checkout:
24
+ - git fetch --tags --unshallow || true
@@ -0,0 +1,138 @@
1
+ Metadata-Version: 2.5
2
+ Name: pytensor_ml
3
+ Version: 0.2.2
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: Documentation, https://pytensor-ml.readthedocs.io
7
+ Project-URL: Source, https://github.com/pymc-devs/pytensor-ml
8
+ Project-URL: Tracker, https://github.com/pymc-devs/pytensor-ml/issues
9
+ Author: Ricardo Vieira
10
+ Author-email: Jesse Grabowski <jessegrabowski@gmail.com>
11
+ License-Expression: Apache-2.0
12
+ License-File: LICENSE
13
+ Keywords: autograd,deep learning,gradient descent,neural network,pymc
14
+ Classifier: Development Status :: 2 - Pre-Alpha
15
+ Classifier: Intended Audience :: Science/Research
16
+ Classifier: Operating System :: MacOS
17
+ Classifier: Operating System :: Microsoft :: Windows
18
+ Classifier: Operating System :: POSIX
19
+ Classifier: Operating System :: Unix
20
+ Classifier: Programming Language :: Python
21
+ Classifier: Programming Language :: Python :: 3
22
+ Classifier: Programming Language :: Python :: 3.12
23
+ Classifier: Programming Language :: Python :: 3.13
24
+ Classifier: Programming Language :: Python :: 3.14
25
+ Classifier: Topic :: Scientific/Engineering
26
+ Requires-Python: >=3.12
27
+ Requires-Dist: numpy
28
+ Requires-Dist: pytensor<4.0.0,>=3.2.3
29
+ Requires-Dist: safetensors
30
+ Provides-Extra: dev
31
+ Requires-Dist: jax; extra == 'dev'
32
+ Requires-Dist: mypy; extra == 'dev'
33
+ Requires-Dist: pre-commit; extra == 'dev'
34
+ Requires-Dist: pytest; extra == 'dev'
35
+ Requires-Dist: pytest-cov; extra == 'dev'
36
+ Requires-Dist: pyyaml; extra == 'dev'
37
+ Requires-Dist: scikit-learn; extra == 'dev'
38
+ Provides-Extra: examples
39
+ Requires-Dist: matplotlib; extra == 'examples'
40
+ Requires-Dist: tqdm; extra == 'examples'
41
+ Description-Content-Type: text/markdown
42
+
43
+ # pytensor_ml
44
+
45
+ A(nother) deep learning library, built on top of [PyTensor](https://github.com/pymc-devs/pytensor).
46
+
47
+ Networks are ordinary PyTensor graphs. You build one out of layers, and everything PyTensor already does —
48
+ symbolic differentiation, graph rewrites, and compilation to Numba, C, JAX, PyTorch, or MLX — applies to it
49
+ unchanged. Training is a compiled function that takes a batch and returns a loss; there is no separate runtime
50
+ or tape.
51
+
52
+ That goes all the way down: layers are graph constructors, parameters are shared variables, and a training
53
+ step is a compiled function whose updates are the optimizer. Because a model is only a graph, it composes with
54
+ any other PyTensor graph — a PyMC model included — as there is nothing else to interoperate with.
55
+
56
+ > **Status: pre-alpha.** The API is still moving, and there is no release-to-release compatibility guarantee yet.
57
+
58
+ ## Installation
59
+
60
+ ```bash
61
+ pip install pytensor-ml
62
+ ```
63
+
64
+ The only hard dependencies are `pytensor`, `numpy`, and `safetensors`. A backend beyond the default (`numba`,
65
+ `jax`, `torch`, `mlx`) is installed separately, and only loads when you actually compile against it.
66
+
67
+ ## Quickstart
68
+
69
+ Train a classifier on scikit-learn's digits, then run inference:
70
+
71
+ ```python
72
+ import numpy as np
73
+ import pytensor
74
+
75
+ pytensor.config.floatX = "float32"
76
+
77
+ from sklearn.datasets import load_digits
78
+
79
+ from pytensor_ml.activations import ReLU
80
+ from pytensor_ml.layers import Input, Linear, Sequential
81
+ from pytensor_ml.loss import CrossEntropy
82
+ from pytensor_ml.model import Model
83
+ from pytensor_ml.optim import adam, chain, clip_by_global_norm, cosine_schedule
84
+ from pytensor_ml.util import DataLoader
85
+
86
+ X, y = load_digits(return_X_y=True)
87
+ X = (X / 16.0).astype("float32")
88
+ y_onehot = np.eye(10, dtype="float32")[y]
89
+
90
+ X_in = Input("X_in", shape=(None, 64))
91
+ network = Sequential(
92
+ Linear("fc1", n_in=64, n_out=128),
93
+ ReLU(),
94
+ Linear("logits", n_in=128, n_out=10),
95
+ )
96
+ model = Model(X_in, network(X_in)).initialize(seed=0)
97
+
98
+ rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
99
+ loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
100
+ step = model.compile_train(rule, loss_fn, ndim_out=2)
101
+
102
+ loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
103
+ for _ in range(500):
104
+ loss_value = step(*loader())
105
+
106
+ accuracy = (model.predict(X).argmax(axis=-1) == y).mean()
107
+ ```
108
+
109
+ `compile_train` builds the loss against a target placeholder, differentiates it, folds in any stateful layer
110
+ updates (batch norm running statistics, RNG advances, the training clock a schedule reads), and compiles a
111
+ one-step function. `predict` compiles a separate inference pass, with dropout removed and batch norm reading
112
+ its running statistics.
113
+
114
+ ## Documentation
115
+
116
+ The full API reference and user guide live at
117
+ [pytensor-ml.readthedocs.io](https://pytensor-ml.readthedocs.io).
118
+
119
+ For worked models end to end — training loops, convolutional and recurrent networks, transformers, saving and
120
+ reloading — see the [examples gallery](https://pytensor-ml.readthedocs.io/en/latest/examples/gallery.html).
121
+
122
+ ## Contributing
123
+
124
+ Contributions are welcome. To get set up:
125
+
126
+ ```bash
127
+ pip install -e ".[dev]"
128
+ pre-commit install
129
+ pytest
130
+ ```
131
+
132
+ Formatting and linting run through `ruff` under pre-commit, and `mypy` checks `pytensor_ml/`; both also run in
133
+ CI. Bug reports and feature requests belong in the
134
+ [issue tracker](https://github.com/pymc-devs/pytensor-ml/issues).
135
+
136
+ ## License
137
+
138
+ Apache 2.0. See [LICENSE](LICENSE).
@@ -0,0 +1,96 @@
1
+ # pytensor_ml
2
+
3
+ A(nother) deep learning library, built on top of [PyTensor](https://github.com/pymc-devs/pytensor).
4
+
5
+ Networks are ordinary PyTensor graphs. You build one out of layers, and everything PyTensor already does —
6
+ symbolic differentiation, graph rewrites, and compilation to Numba, C, JAX, PyTorch, or MLX — applies to it
7
+ unchanged. Training is a compiled function that takes a batch and returns a loss; there is no separate runtime
8
+ or tape.
9
+
10
+ That goes all the way down: layers are graph constructors, parameters are shared variables, and a training
11
+ step is a compiled function whose updates are the optimizer. Because a model is only a graph, it composes with
12
+ any other PyTensor graph — a PyMC model included — as there is nothing else to interoperate with.
13
+
14
+ > **Status: pre-alpha.** The API is still moving, and there is no release-to-release compatibility guarantee yet.
15
+
16
+ ## Installation
17
+
18
+ ```bash
19
+ pip install pytensor-ml
20
+ ```
21
+
22
+ The only hard dependencies are `pytensor`, `numpy`, and `safetensors`. A backend beyond the default (`numba`,
23
+ `jax`, `torch`, `mlx`) is installed separately, and only loads when you actually compile against it.
24
+
25
+ ## Quickstart
26
+
27
+ Train a classifier on scikit-learn's digits, then run inference:
28
+
29
+ ```python
30
+ import numpy as np
31
+ import pytensor
32
+
33
+ pytensor.config.floatX = "float32"
34
+
35
+ from sklearn.datasets import load_digits
36
+
37
+ from pytensor_ml.activations import ReLU
38
+ from pytensor_ml.layers import Input, Linear, Sequential
39
+ from pytensor_ml.loss import CrossEntropy
40
+ from pytensor_ml.model import Model
41
+ from pytensor_ml.optim import adam, chain, clip_by_global_norm, cosine_schedule
42
+ from pytensor_ml.util import DataLoader
43
+
44
+ X, y = load_digits(return_X_y=True)
45
+ X = (X / 16.0).astype("float32")
46
+ y_onehot = np.eye(10, dtype="float32")[y]
47
+
48
+ X_in = Input("X_in", shape=(None, 64))
49
+ network = Sequential(
50
+ Linear("fc1", n_in=64, n_out=128),
51
+ ReLU(),
52
+ Linear("logits", n_in=128, n_out=10),
53
+ )
54
+ model = Model(X_in, network(X_in)).initialize(seed=0)
55
+
56
+ rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
57
+ loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
58
+ step = model.compile_train(rule, loss_fn, ndim_out=2)
59
+
60
+ loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
61
+ for _ in range(500):
62
+ loss_value = step(*loader())
63
+
64
+ accuracy = (model.predict(X).argmax(axis=-1) == y).mean()
65
+ ```
66
+
67
+ `compile_train` builds the loss against a target placeholder, differentiates it, folds in any stateful layer
68
+ updates (batch norm running statistics, RNG advances, the training clock a schedule reads), and compiles a
69
+ one-step function. `predict` compiles a separate inference pass, with dropout removed and batch norm reading
70
+ its running statistics.
71
+
72
+ ## Documentation
73
+
74
+ The full API reference and user guide live at
75
+ [pytensor-ml.readthedocs.io](https://pytensor-ml.readthedocs.io).
76
+
77
+ For worked models end to end — training loops, convolutional and recurrent networks, transformers, saving and
78
+ reloading — see the [examples gallery](https://pytensor-ml.readthedocs.io/en/latest/examples/gallery.html).
79
+
80
+ ## Contributing
81
+
82
+ Contributions are welcome. To get set up:
83
+
84
+ ```bash
85
+ pip install -e ".[dev]"
86
+ pre-commit install
87
+ pytest
88
+ ```
89
+
90
+ Formatting and linting run through `ruff` under pre-commit, and `mypy` checks `pytensor_ml/`; both also run in
91
+ CI. Bug reports and feature requests belong in the
92
+ [issue tracker](https://github.com/pymc-devs/pytensor-ml/issues).
93
+
94
+ ## License
95
+
96
+ Apache 2.0. See [LICENSE](LICENSE).
@@ -0,0 +1,35 @@
1
+ # The environment Read the Docs builds from. `pixi run docs-build` uses the equivalent list under
2
+ # [tool.pixi.feature.docs] in pyproject.toml; the two have to move together.
3
+ name: pytensor_ml-docs
4
+ channels:
5
+ - conda-forge
6
+ - nodefaults
7
+
8
+ dependencies:
9
+ - python>=3.12
10
+ # Runtime deps: autodoc imports pytensor_ml, so the full runtime stack has to be in scope.
11
+ - pytensor>=3.3.0,<3.4.0
12
+ - numpy
13
+ - safetensors
14
+ # The gallery extension renders notebook thumbnails with matplotlib.
15
+ - matplotlib
16
+
17
+ # Docs build deps.
18
+ - ipython
19
+ - jupyter
20
+ - sphinx>=7
21
+ - pydata-sphinx-theme
22
+ - myst-nb
23
+ - numpydoc
24
+ - sphinx-copybutton
25
+ - sphinx-design
26
+ - sphinx-codeautolink
27
+ - sphinx-sitemap
28
+ - sphinx-notfound-page
29
+ - sphinx-autobuild
30
+ - jupyter-sphinx
31
+ - sphinxcontrib-bibtex
32
+ - pip
33
+ - pip:
34
+ # pytensor_ml itself so autodoc resolves the current source tree.
35
+ - -e ..
@@ -20,13 +20,6 @@ dependencies:
20
20
  - pytest-mock
21
21
  - pyyaml # tests/test_workflow_groups.py reads the CI matrix
22
22
 
23
- # For building docs
24
- - sphinx
25
- - sphinx_rtd_theme
26
- - pygments
27
- - pydot
28
- - ipython
29
-
30
23
  # developer tools
31
24
  - pre-commit
32
25
  # Pinned rather than floated: a local run that disagrees with CI about the version is a local run
@@ -0,0 +1,16 @@
1
+ # Sphinx build output
2
+ build/
3
+ jupyter_execute/
4
+
5
+ # Python cache
6
+ __pycache__/
7
+ *.pyc
8
+
9
+ # Autosummary-generated stubs (one .rst per object, regenerated on build)
10
+ source/api/generated/
11
+ source/api/**/generated/
12
+ source/api/**/classmethods/
13
+
14
+ # Notebook gallery artifacts written by docs/sphinxext/generate_gallery.py
15
+ source/examples/
16
+ source/_thumbnails/
@@ -0,0 +1,34 @@
1
+ # Minimal Makefile for Sphinx documentation.
2
+ # Mirrors the standard sphinx-quickstart output.
3
+
4
+ SPHINXOPTS ?=
5
+ SPHINXBUILD ?= sphinx-build
6
+ SOURCEDIR = source
7
+ BUILDDIR = build
8
+
9
+ .PHONY: help clean html show livehtml linkcheck Makefile
10
+
11
+ help:
12
+ @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
13
+
14
+ clean:
15
+ rm -rf "$(BUILDDIR)"
16
+ rm -rf "$(SOURCEDIR)"/api/generated
17
+ rm -rf "$(SOURCEDIR)"/api/*/generated
18
+ rm -rf "$(SOURCEDIR)"/_thumbnails
19
+ rm -f "$(SOURCEDIR)"/examples/gallery.rst
20
+
21
+ html:
22
+ @$(SPHINXBUILD) -M html "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
23
+
24
+ show: html
25
+ @open "$(BUILDDIR)/html/index.html"
26
+
27
+ livehtml:
28
+ sphinx-autobuild "$(SOURCEDIR)" "$(BUILDDIR)/html" $(SPHINXOPTS) $(O)
29
+
30
+ linkcheck:
31
+ @$(SPHINXBUILD) -M linkcheck "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
32
+
33
+ %: Makefile
34
+ @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
@@ -0,0 +1,29 @@
1
+ @ECHO OFF
2
+
3
+ pushd %~dp0
4
+
5
+ REM Command file for Sphinx documentation
6
+
7
+ if "%SPHINXBUILD%" == "" (
8
+ set SPHINXBUILD=sphinx-build
9
+ )
10
+ set SOURCEDIR=source
11
+ set BUILDDIR=build
12
+
13
+ %SPHINXBUILD% >NUL 2>NUL
14
+ if errorlevel 9009 (
15
+ echo.
16
+ echo.The 'sphinx-build' command was not found. Install Sphinx, then re-run.
17
+ exit /b 1
18
+ )
19
+
20
+ if "%1" == "" goto help
21
+
22
+ %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
23
+ goto end
24
+
25
+ :help
26
+ %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
27
+
28
+ :end
29
+ popd
@@ -0,0 +1,34 @@
1
+ {{ fullname | escape | underline}}
2
+
3
+ .. currentmodule:: {{ module }}
4
+
5
+ .. autoclass:: {{ objname }}
6
+
7
+ {% block methods %}
8
+ {% if methods %}
9
+
10
+ .. rubric:: Methods
11
+
12
+ .. autosummary::
13
+ :toctree: classmethods
14
+
15
+ {% for item in methods %}
16
+ {%- if item not in inherited_members %}
17
+ {{ objname }}.{{ item }}
18
+ {% endif %}
19
+ {%- endfor %}
20
+ {% endif %}
21
+ {% endblock %}
22
+
23
+ {% block attributes %}
24
+ {% if attributes %}
25
+ .. rubric:: Attributes
26
+
27
+ .. autosummary::
28
+ {% for item in attributes %}
29
+ {%- if item not in inherited_members %}
30
+ ~{{ name }}.{{ item }}
31
+ {% endif %}
32
+ {%- endfor %}
33
+ {% endif %}
34
+ {% endblock %}
@@ -0,0 +1,17 @@
1
+ Activations
2
+ ===========
3
+
4
+ .. currentmodule:: pytensor_ml.activations
5
+
6
+ .. autosummary::
7
+ :toctree: generated/
8
+
9
+ Activation
10
+ ReLU
11
+ LeakyReLU
12
+ GELU
13
+ Swish
14
+ Sigmoid
15
+ SoftPlus
16
+ Softmax
17
+ Tanh
@@ -0,0 +1,101 @@
1
+ Layers
2
+ ======
3
+
4
+ .. currentmodule:: pytensor_ml.layers
5
+
6
+ Base
7
+ ----
8
+
9
+ .. autosummary::
10
+ :toctree: generated/
11
+
12
+ Layer
13
+
14
+ Combinators
15
+ -----------
16
+
17
+ .. autosummary::
18
+ :toctree: generated/
19
+
20
+ Input
21
+ Sequential
22
+ Concatenate
23
+ Flatten
24
+ Squeeze
25
+
26
+ Dense
27
+ -----
28
+
29
+ .. autosummary::
30
+ :toctree: generated/
31
+
32
+ Linear
33
+ Embedding
34
+
35
+ Convolution and pooling
36
+ -----------------------
37
+
38
+ .. autosummary::
39
+ :toctree: generated/
40
+
41
+ Conv1D
42
+ Conv2D
43
+ ConvTranspose1D
44
+ ConvTranspose2D
45
+ MaxPool1D
46
+ MaxPool2D
47
+ AvgPool1D
48
+ AvgPool2D
49
+
50
+ Padding
51
+ -------
52
+
53
+ .. autosummary::
54
+ :toctree: generated/
55
+
56
+ ZeroPad1D
57
+ ZeroPad2D
58
+ ConstantPad1D
59
+ ConstantPad2D
60
+ ReflectionPad1D
61
+ ReflectionPad2D
62
+ ReplicationPad1D
63
+ ReplicationPad2D
64
+
65
+ Normalization and regularization
66
+ --------------------------------
67
+
68
+ .. autosummary::
69
+ :toctree: generated/
70
+
71
+ BatchNorm
72
+ LayerNorm
73
+ Dropout
74
+
75
+ Recurrent
76
+ ---------
77
+
78
+ .. autosummary::
79
+ :toctree: generated/
80
+
81
+ RNN
82
+ LSTM
83
+ GRU
84
+ Bidirectional
85
+ Recurrent
86
+ RecurrentCell
87
+ ElmanCell
88
+ LSTMCell
89
+ GRUCell
90
+
91
+ Attention and transformers
92
+ --------------------------
93
+
94
+ .. autosummary::
95
+ :toctree: generated/
96
+
97
+ MultiheadAttention
98
+ CausalSelfAttention
99
+ FeedForward
100
+ TransformerBlock
101
+ scaled_dot_product_attention
@@ -0,0 +1,12 @@
1
+ Losses
2
+ ======
3
+
4
+ .. currentmodule:: pytensor_ml.loss
5
+
6
+ .. autosummary::
7
+ :toctree: generated/
8
+
9
+ Loss
10
+ SquaredError
11
+ CrossEntropy
12
+ supervised_loss
@@ -0,0 +1,9 @@
1
+ Model
2
+ =====
3
+
4
+ .. currentmodule:: pytensor_ml.model
5
+
6
+ .. autosummary::
7
+ :toctree: generated/
8
+
9
+ Model