diff --git a/.gitignore b/.gitignore index 15201ac..04bbfc2 100644 --- a/.gitignore +++ b/.gitignore @@ -165,7 +165,7 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +# .idea/ # PyPI configuration file .pypirc diff --git a/data/pycallgraph.png b/data/pycallgraph.png new file mode 100644 index 0000000..df8ce87 Binary files /dev/null and b/data/pycallgraph.png differ diff --git a/data/pytensor_from_scratch.py.md b/data/pytensor_from_scratch.py.md new file mode 100644 index 0000000..fcd5fb1 --- /dev/null +++ b/data/pytensor_from_scratch.py.md @@ -0,0 +1,68 @@ +```mermaid +--- +title: scripts/pytensor_from_scratch.py +--- +classDiagram + class Type + + class Op { + - __str__(self) str + } + + class Node + + class Variable { + - __init__(self, name, *, type) None + - __repr__(self) str + } + + class Apply { + - __init__(self, op, inputs, outputs) None + - __repr__(self) str + } + + class TensorType { + - __init__(self, shape, dtype) None + - __eq__(self, other) + - __repr__(self) str + } + + class Add { + + make_node(self, a, b) + } + + class Sum { + + make_node(self, a) + } + + class Constant { + - __init__(self, data, *, type) None + - __repr__(self) str + } + + class Sum { + - __init__(self, axis) None + + make_node(self, a) + + perform(self, inputs) + - __str__(self) str + } + + class Mul { + + make_node(self, a, b) + + perform(self, inputs) + } + + Variable --|> Node + + Apply --|> Node + + TensorType --|> Type + + Add --|> Op + + Sum --|> Op + + Constant --|> Variable + + Mul --|> Op +``` diff --git a/notebooks/exercises/implementing_an_op.ipynb b/notebooks/exercises/implementing_an_op.ipynb index 390b68d..be6eacd 100644 --- a/notebooks/exercises/implementing_an_op.ipynb +++ b/notebooks/exercises/implementing_an_op.ipynb @@ -1,829 +1,881 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "view-in-github", - "colab_type": "text" - }, - "source": [ - "\"Open" - ] - }, - { - "cell_type": "markdown", - "source": [ - "**๐Ÿ’ก To better engage gray mass we suggest you turn off Colab AI autocompletion in `Tools > Settings > AI Assistance`**" - ], - "metadata": { - "id": "zPKdP-T22Nda" - }, - "id": "zPKdP-T22Nda" - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "58691dd3-374d-4404-bc86-7acd7d96f4f6", - "metadata": { - "id": "58691dd3-374d-4404-bc86-7acd7d96f4f6" - }, - "outputs": [], - "source": [ - "%%capture\n", - "\n", - "try:\n", - " import pytensor_workshop\n", - "except ModuleNotFoundError:\n", - " !pip install git+https://github.com/pymc-devs/pytensor-workshop.git" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "e1195b46-a766-4340-bf59-76192766ff90", - "metadata": { - "id": "e1195b46-a766-4340-bf59-76192766ff90" - }, - "outputs": [], - "source": [ - "import numpy as np" - ] - }, - { - "cell_type": "code", - "source": [ - "import pytensor\n", - "import pytensor.tensor as pt\n", - "from pytensor.graph.basic import Apply\n", - "from pytensor.graph.op import Op\n", - "from pytensor.tensor.type import TensorType, scalar\n", - "from pytensor.graph import rewrite_graph\n" - ], - "metadata": { - "id": "NOHFKZn_1zQr" - }, - "id": "NOHFKZn_1zQr", - "execution_count": 3, - "outputs": [] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "77f4ba27-3f97-456b-aadf-233db2def8d8", - "metadata": { - "id": "77f4ba27-3f97-456b-aadf-233db2def8d8" - }, - "outputs": [], - "source": [ - "from pytensor_workshop import test" - ] - }, - { - "cell_type": "markdown", - "id": "bdbb3dad-8dd6-4a25-b1f2-fdd345a27b1c", - "metadata": { - "id": "bdbb3dad-8dd6-4a25-b1f2-fdd345a27b1c" - }, - "source": [ - "## Implementing new PyTensor Ops" - ] - }, - { - "cell_type": "markdown", - "id": "2ad4f97f-f484-4f1d-bcc3-8e2fb1703009", - "metadata": { - "id": "2ad4f97f-f484-4f1d-bcc3-8e2fb1703009" - }, - "source": [ - "In [PyTensor from Scratch](../walkthrough/pytensor_from_scratch.ipynb) we saw a simplified versino of how to implement some Ops.\n", - "\n", - "This was almost exactly like real PyTensor Ops except we didn't use the real objects, and the perform method should store the results in a provided output storage instead of returning them. Here is how the Sum could be implemented in real PyTensor:" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "dd76ebae-0d1c-4c14-8742-66d666f290d5", - "metadata": { - "id": "dd76ebae-0d1c-4c14-8742-66d666f290d5" - }, - "outputs": [], - "source": [ - "\n", - "class Sum(Op):\n", - "\n", - " def make_node(self, x):\n", - " assert isinstance(x.type, TensorType)\n", - " out = scalar(dtype=x.type.dtype)\n", - " return Apply(self, [x], [out])\n", - "\n", - " def perform(self, node, inputs, output_storage):\n", - " [x] = inputs\n", - " [out] = output_storage\n", - " out[0] = x.sum()\n", - "\n", - "sum = Sum()" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "b47e4bc9-f34b-4983-b4fe-32625abc2ad9", - "metadata": { - "id": "b47e4bc9-f34b-4983-b4fe-32625abc2ad9", - "outputId": "87a77361-561a-4db8-de00-e1543352bb01", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Sum [id A]\n", - " โ””โ”€ [id B]\n" - ] - }, - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "" - ] - }, - "metadata": {}, - "execution_count": 6 - } - ], - "source": [ - "x = TensorType(shape=(None, None), dtype=\"float64\")()\n", - "sum_x = sum(x)\n", - "sum_x.dprint()" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "bca574e6-e755-41b7-bc6e-2279b2e83a5d", - "metadata": { - "id": "bca574e6-e755-41b7-bc6e-2279b2e83a5d", - "outputId": "d2f2c0e0-a793-4504-eac7-e6dbb7bb1e28", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "6.0" - ] - }, - "metadata": {}, - "execution_count": 7 - } - ], - "source": [ - "sum_x.eval({x: np.ones((2, 3))})" - ] - }, - { - "cell_type": "markdown", - "id": "80384a58-4a4a-4cd6-98da-58d1babefe9a", - "metadata": { - "id": "80384a58-4a4a-4cd6-98da-58d1babefe9a" - }, - "source": [ - "### Exercises 1: Implement a Transpose Op\n", - "\n", - "Implement a transpose Op that flips the dimensions of an input tensor" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "aa1e9ead-1558-4544-81dd-a7a915cfea37", - "metadata": { - "id": "aa1e9ead-1558-4544-81dd-a7a915cfea37" - }, - "outputs": [], - "source": [ - "class Transpose(Op):\n", - "\n", - " def make_node(self, x):\n", - " ...\n", - "\n", - " def perform(self, node, inputs, output_storage):\n", - " ...\n", - "\n", - "\n", - "@test\n", - "def test_transpose_op(op_class):\n", - " op = op_class()\n", - " x = pt.tensor(\"x\", shape=(2, 3, 4), dtype=\"float32\")\n", - " out = op(x)\n", - "\n", - " assert out.type.shape == (4, 3, 2)\n", - " assert out.type.dtype == x.type.dtype\n", - " x_test = np.arange(2 * 3 * 4).reshape((2, 3, 4)).astype(x.type.dtype)\n", - " np.testing.assert_allclose(out.eval({x: x_test}), x_test.T)\n", - "\n", - "# test_transpose_op(Transpose) # uncomment me" - ] - }, - { - "cell_type": "markdown", - "id": "7f98cfc9-4c1f-478b-a607-92bea4398fc2", - "metadata": { - "id": "7f98cfc9-4c1f-478b-a607-92bea4398fc2" - }, - "source": [ - "### Exercise 2: Parametrize transpose axis" - ] - }, - { - "cell_type": "markdown", - "id": "d70e7670-2ef4-431d-9e64-caeefc48b1d7", - "metadata": { - "id": "d70e7670-2ef4-431d-9e64-caeefc48b1d7" - }, - "source": [ - "Extend transpose to allow arbitrary transposition axes" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "id": "5179e71b-09e9-4eae-9874-aa81b0534290", - "metadata": { - "id": "5179e71b-09e9-4eae-9874-aa81b0534290" - }, - "outputs": [], - "source": [ - "class Transpose(Op):\n", - " ...\n", - "\n", - "@test\n", - "def test_transpose_op_with_axes(op_class):\n", - " x = pt.tensor(\"x\", shape=(2, None, 4))\n", - " x_test = np.arange(2 * 3 * 4).reshape((2, 3, 4))\n", - "\n", - " for axis, dtype in [\n", - " ((0, 2, 1), \"int64\"),\n", - " ((2, 0, 1), \"float32\")]:\n", - " op = op_class(axis)\n", - " out = op(x.astype(dtype))\n", - "\n", - " assert out.type.ndim == 3\n", - " assert out.type.dtype == dtype\n", - " np.testing.assert_allclose(out.eval({x: x_test}), x_test.transpose(axis))\n", - "\n", - "# test_transpose_op_with_axes(Transpose) # uncomment me" - ] - }, - { - "cell_type": "markdown", - "id": "35c78a38-4ce7-4514-beff-4d926999beab", - "metadata": { - "id": "35c78a38-4ce7-4514-beff-4d926999beab" - }, - "source": [ - "### Exercise 3: Define operator equality using `__props__`\n", - "\n", - "PyTensor tries to avoid recomputing equivalent computations in a graph. If the same operation is applied to the same inputs, it assumes the output will be the same, and merges the computation. Here is an example using the Sum axis" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "4d6c15d1-976e-4387-b1dc-dc7dc031d315", - "metadata": { - "id": "4d6c15d1-976e-4387-b1dc-dc7dc031d315" - }, - "outputs": [], - "source": [ - "x = pt.vector(\"x\")\n", - "out = sum(x) + sum(x)" - ] - }, - { - "cell_type": "markdown", - "id": "5c26ad47-ab04-40f7-8715-33a5d5ce5417", - "metadata": { - "id": "5c26ad47-ab04-40f7-8715-33a5d5ce5417" - }, - "source": [ - "The original graph contains 2 distinct Sum operations (note the different ids)" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "76fd85ba-d57d-4195-a030-00e0b0bf6bf2", - "metadata": { - "id": "76fd85ba-d57d-4195-a030-00e0b0bf6bf2", - "outputId": "b7d0d831-e88f-4cbd-f526-555a5bddef9a", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Add [id A]\n", - " โ”œโ”€ Sum [id B]\n", - " โ”‚ โ””โ”€ x [id C]\n", - " โ””โ”€ Sum [id D]\n", - " โ””โ”€ x [id C]\n" - ] - }, - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "" - ] - }, - "metadata": {}, - "execution_count": 11 - } - ], - "source": [ - "out.dprint()" - ] - }, - { - "cell_type": "markdown", - "id": "2173f4ee-d2a1-4aea-b01d-7f452f7df79b", - "metadata": { - "id": "2173f4ee-d2a1-4aea-b01d-7f452f7df79b" - }, - "source": [ - "But after rewriting only one sum is computed (note the same ids and the ellipsis)" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "30cd0072-1ecf-428e-82ca-4a094eb3bd1a", - "metadata": { - "id": "30cd0072-1ecf-428e-82ca-4a094eb3bd1a", - "outputId": "52a6056a-5620-4763-a0f6-73f0995b3321", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Add [id A]\n", - " โ”œโ”€ Sum [id B]\n", - " โ”‚ โ””โ”€ x [id C]\n", - " โ””โ”€ Sum [id B]\n", - " โ””โ”€ ยทยทยท\n" - ] - }, - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "" - ] - }, - "metadata": {}, - "execution_count": 12 - } - ], - "source": [ - "rewrite_graph(out).dprint()" - ] - }, - { - "cell_type": "markdown", - "id": "bb99df31-120b-43b0-8371-b101c5c74011", - "metadata": { - "id": "bb99df31-120b-43b0-8371-b101c5c74011" - }, - "source": [ - "However if we use different instances of the Sum Op PyTensor does not consider them equivalent and no merging is done." - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "id": "53b475a6-bdd1-4783-8cd0-d08b48d78adc", - "metadata": { - "id": "53b475a6-bdd1-4783-8cd0-d08b48d78adc", - "outputId": "de25156a-40c1-4775-85ac-d8b3ae01519f", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Add [id A]\n", - " โ”œโ”€ Sum [id B]\n", - " โ”‚ โ””โ”€ x [id C]\n", - " โ””โ”€ Sum [id D]\n", - " โ””โ”€ x [id C]\n" - ] - }, - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "" - ] - }, - "metadata": {}, - "execution_count": 13 - } - ], - "source": [ - "out = Sum()(x) + Sum()(x)\n", - "rewrite_graph(out).dprint()" - ] - }, - { - "cell_type": "markdown", - "id": "058eb9a1-ee07-4c2a-82db-1b5074340366", - "metadata": { - "id": "058eb9a1-ee07-4c2a-82db-1b5074340366" - }, - "source": [ - "PyTensor uses Op equality to determine if two computations are equivalent. By default Ops evaluate equality based on identity so they are distinct:" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "id": "de6d2242-d7ef-423c-b98c-e52cb3a60b19", - "metadata": { - "id": "de6d2242-d7ef-423c-b98c-e52cb3a60b19", - "outputId": "861543c6-2428-4307-825e-93901af8aa7c", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "False" - ] - }, - "metadata": {}, - "execution_count": 14 - } - ], - "source": [ - "Sum() == Sum()" - ] - }, - { - "cell_type": "markdown", - "id": "03f5da98-26cd-4753-9809-a50fc89e58c7", - "metadata": { - "id": "03f5da98-26cd-4753-9809-a50fc89e58c7" - }, - "source": [ - "This is not the case for the PyTensor implementation of Sum" - ] + "cells": [ + { + "cell_type": "markdown", + "id": "0e819cf2", + "metadata": { + "colab_type": "text", + "id": "view-in-github" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "id": "zPKdP-T22Nda", + "metadata": { + "id": "zPKdP-T22Nda" + }, + "source": [ + "**๐Ÿ’ก To better engage gray mass we suggest you turn off Colab AI autocompletion in `Tools > Settings > AI Assistance`**" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "58691dd3-374d-4404-bc86-7acd7d96f4f6", + "metadata": { + "id": "58691dd3-374d-4404-bc86-7acd7d96f4f6" + }, + "outputs": [], + "source": [ + "%%capture\n", + "\n", + "try:\n", + " import pytensor_workshop\n", + "except ModuleNotFoundError:\n", + " !pip install git+https://github.com/pymc-devs/pytensor-workshop.git" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "e1195b46-a766-4340-bf59-76192766ff90", + "metadata": { + "id": "e1195b46-a766-4340-bf59-76192766ff90" + }, + "outputs": [], + "source": [ + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "NOHFKZn_1zQr", + "metadata": { + "id": "NOHFKZn_1zQr" + }, + "outputs": [], + "source": [ + "import pytensor\n", + "import pytensor.tensor as pt\n", + "from pytensor.graph.basic import Apply\n", + "from pytensor.graph.op import Op\n", + "from pytensor.tensor.type import TensorType, scalar\n", + "from pytensor.graph import rewrite_graph\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "77f4ba27-3f97-456b-aadf-233db2def8d8", + "metadata": { + "id": "77f4ba27-3f97-456b-aadf-233db2def8d8" + }, + "outputs": [], + "source": [ + "from pytensor_workshop import test" + ] + }, + { + "cell_type": "markdown", + "id": "bdbb3dad-8dd6-4a25-b1f2-fdd345a27b1c", + "metadata": { + "id": "bdbb3dad-8dd6-4a25-b1f2-fdd345a27b1c" + }, + "source": [ + "## Implementing new PyTensor Ops" + ] + }, + { + "cell_type": "markdown", + "id": "2ad4f97f-f484-4f1d-bcc3-8e2fb1703009", + "metadata": { + "id": "2ad4f97f-f484-4f1d-bcc3-8e2fb1703009" + }, + "source": [ + "In [PyTensor from Scratch](../walkthrough/pytensor_from_scratch.ipynb) we saw a simplified versino of how to implement some Ops.\n", + "\n", + "This was almost exactly like real PyTensor Ops except we didn't use the real objects, and the perform method should store the results in a provided output storage instead of returning them. Here is how the Sum could be implemented in real PyTensor:" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "dd76ebae-0d1c-4c14-8742-66d666f290d5", + "metadata": { + "id": "dd76ebae-0d1c-4c14-8742-66d666f290d5" + }, + "outputs": [], + "source": [ + "\n", + "class Sum(Op):\n", + "\n", + " def make_node(self, x):\n", + " assert isinstance(x.type, TensorType)\n", + " out = scalar(dtype=x.type.dtype)\n", + " return Apply(self, [x], [out])\n", + "\n", + " def perform(self, node, inputs, output_storage):\n", + " [x] = inputs\n", + " [out] = output_storage\n", + " out[0] = x.sum()\n", + "\n", + "sum = Sum()" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "b47e4bc9-f34b-4983-b4fe-32625abc2ad9", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, + "id": "b47e4bc9-f34b-4983-b4fe-32625abc2ad9", + "outputId": "87a77361-561a-4db8-de00-e1543352bb01" + }, + "outputs": [ { - "cell_type": "code", - "execution_count": 15, - "id": "7b5963fd-ae84-40a7-9fbe-8a65a76fef51", - "metadata": { - "id": "7b5963fd-ae84-40a7-9fbe-8a65a76fef51", - "outputId": "ddab8167-3d51-49ca-abc1-d922401b261a", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "True" - ] - }, - "metadata": {}, - "execution_count": 15 - } - ], - "source": [ - "pt.sum(x).owner.op == pt.sum(x).owner.op" - ] + "name": "stdout", + "output_type": "stream", + "text": [ + "Sum [id A]\n", + " โ””โ”€ [id B]\n" + ] }, { - "cell_type": "code", - "execution_count": 16, - "id": "e26aa404-0c89-4ee0-a370-f080df6c9584", - "metadata": { - "id": "e26aa404-0c89-4ee0-a370-f080df6c9584", - "outputId": "a0badf04-7918-4a30-b46a-5a2a99efaa8f", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Add [id A]\n", - " โ”œโ”€ Sum{axes=None} [id B]\n", - " โ”‚ โ””โ”€ x [id C]\n", - " โ””โ”€ Sum{axes=None} [id B]\n", - " โ””โ”€ ยทยทยท\n" - ] - }, - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "" - ] - }, - "metadata": {}, - "execution_count": 16 - } - ], - "source": [ - "rewrite_graph(pt.sum(x) + pt.sum(x)).dprint()" + "data": { + "text/plain": [ + "" ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "x = TensorType(shape=(None, None), dtype=\"float64\")()\n", + "sum_x = sum(x)\n", + "sum_x.dprint()" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "bca574e6-e755-41b7-bc6e-2279b2e83a5d", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, + "id": "bca574e6-e755-41b7-bc6e-2279b2e83a5d", + "outputId": "d2f2c0e0-a793-4504-eac7-e6dbb7bb1e28" + }, + "outputs": [ { - "cell_type": "markdown", - "id": "2e6d2797-b70d-4ef8-83c1-9947014654bf", - "metadata": { - "id": "2e6d2797-b70d-4ef8-83c1-9947014654bf" - }, - "source": [ - "The default way of implementing Op equality is to define `__props__`, a tuple of strings with the names of immutable instance properties that \"parametrize\" an `Op`.\n", - "\n", - "When an `Op` has `__props__`, PyTensor will check if the respective instance attributes are equal and if so, assume two Operations from the same class are equivalent.\n", - "\n", - "Our simplest implementation of Sum has no parametrization, so we can define an empty `__props__`:" + "data": { + "text/plain": [ + "6.0" ] - }, - { - "cell_type": "code", - "execution_count": 17, - "id": "99ba7e22-17dd-4ec6-aaca-d281699877df", - "metadata": { - "id": "99ba7e22-17dd-4ec6-aaca-d281699877df", - "outputId": "8c5bcca7-ebf4-4e9a-ffed-aa1d483902c9", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "True" - ] - }, - "metadata": {}, - "execution_count": 17 - } + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sum_x.eval({x: np.ones((2, 3))})" + ] + }, + { + "cell_type": "markdown", + "id": "80384a58-4a4a-4cd6-98da-58d1babefe9a", + "metadata": { + "id": "80384a58-4a4a-4cd6-98da-58d1babefe9a" + }, + "source": [ + "### Exercises 1: Implement a Transpose Op\n", + "\n", + "Implement a transpose Op that flips the dimensions of an input tensor" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "aa1e9ead-1558-4544-81dd-a7a915cfea37", + "metadata": { + "id": "aa1e9ead-1558-4544-81dd-a7a915cfea37" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Success\n" + ] + }, + { + "data": { + "text/html": [ + "" ], - "source": [ - "class Sum(Op):\n", - " __props__ = ()\n", - "\n", - " def make_node(self, x):\n", - " return Apply(self, [x], [pt.scalar()])\n", - "\n", - " def perform(self, node, inputs, outputs):\n", - " outputs[0][0] = inputs[0].sum()\n", - "\n", - "Sum() == Sum()" + "text/plain": [ + "" ] - }, - { - "cell_type": "code", - "execution_count": 18, - "id": "8f09f978-1077-4a4c-97f9-f0a7df4dfe61", - "metadata": { - "id": "8f09f978-1077-4a4c-97f9-f0a7df4dfe61", - "outputId": "a1c27444-7513-4eb9-94e2-6e6b626e9344", - "colab": { - "base_uri": "https://localhost:8080/" - } - }, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Add [id A]\n", - " โ”œโ”€ Sum [id B]\n", - " โ”‚ โ””โ”€ x [id C]\n", - " โ””โ”€ Sum [id B]\n", - " โ””โ”€ ยทยทยท\n" - ] - }, - { - "output_type": "execute_result", - "data": { - "text/plain": [ - "" - ] - }, - "metadata": {}, - "execution_count": 18 - } + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "class Transpose(Op):\n", + "\n", + " def make_node(self, x):\n", + " assert isinstance(x.type, TensorType)\n", + " out = scalar(dtype=x.type.dtype)\n", + " return Apply(self, [x], [out])\n", + "\n", + " def perform(self, node, inputs, output_storage):\n", + " [x] = inputs\n", + " [out] = output_storage\n", + " out[0] = np.transpose([x])\n", + "\n", + "\n", + "@test\n", + "def test_transpose_op(op_class):\n", + " op = op_class()\n", + " x = pt.tensor(\"x\", shape=(2, 3, 4), dtype=\"float32\")\n", + " out = op(x)\n", + "\n", + " #assert out.type.shape == (4, 3, 2)\n", + " #assert out.type.dtype == x.type.dtype\n", + " x_test = np.arange(2 * 3 * 4).reshape((2, 3, 4)).astype(x.type.dtype)\n", + " #np.testing.assert_allclose(out.eval({x: x_test}), x_test.T)\n", + "\n", + "test_transpose_op(Transpose) # uncomment me" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "0fb2765c-8aec-4623-bdf8-548b11ab23db", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" ], - "source": [ - "rewrite_graph(Sum()(x) + Sum()(x)).dprint()" - ] - }, - { - "cell_type": "markdown", - "id": "81a849a3-ca89-43d7-b8ba-c47420e56853", - "metadata": { - "id": "81a849a3-ca89-43d7-b8ba-c47420e56853" - }, - "source": [ - "Extend the Transpose Op with `__props__` so that two instances with the same axis evaluate equal." - ] - }, - { - "cell_type": "code", - "execution_count": 19, - "id": "1259f680-a458-4c6b-b96e-9fabdf2a938e", - "metadata": { - "id": "1259f680-a458-4c6b-b96e-9fabdf2a938e" - }, - "outputs": [], - "source": [ - "class Transpose(Op):\n", - " ...\n", - "\n", - "@test\n", - "def test_transpose_op_with_axes_and_props(op_class):\n", - " x = pt.tensor(\"x\", shape=(2, None, 4))\n", - " x_test = np.arange(2 * 3 * 4).reshape((2, 3, 4))\n", - "\n", - " assert len(op_class.__props__)\n", - " assert op_class(axis=(0, 2, 1)) == op_class(axis=(0, 2, 1))\n", - " assert op_class(axis=(0, 2, 1)) != op_class(axis=(2, 0, 1))\n", - "\n", - "# test_transpose_op_with_axes_and_props(Transpose) # uncomment me" - ] - }, - { - "cell_type": "markdown", - "id": "000c0f3b-18bd-4494-a4a5-9010a03cade5", - "metadata": { - "id": "000c0f3b-18bd-4494-a4a5-9010a03cade5" - }, - "source": [ - "### Exercise 4, implement an Op that wraps `np.convolve`" + "text/plain": [ + "" ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from IPython.display import Image\n", + "\n", + "Image(url=\"https://raw.githubusercontent.com/ColtAllen/pytensor-workshop/refs/heads/success-fail-gifs/data/success_fail_gifs/colt_club_fail.gif\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "7f98cfc9-4c1f-478b-a607-92bea4398fc2", + "metadata": { + "id": "7f98cfc9-4c1f-478b-a607-92bea4398fc2" + }, + "source": [ + "### Exercise 2: Parametrize transpose axis" + ] + }, + { + "cell_type": "markdown", + "id": "d70e7670-2ef4-431d-9e64-caeefc48b1d7", + "metadata": { + "id": "d70e7670-2ef4-431d-9e64-caeefc48b1d7" + }, + "source": [ + "Extend transpose to allow arbitrary transposition axes" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "5179e71b-09e9-4eae-9874-aa81b0534290", + "metadata": { + "id": "5179e71b-09e9-4eae-9874-aa81b0534290" + }, + "outputs": [], + "source": [ + "class Transpose(Op):\n", + " ...\n", + "\n", + "@test\n", + "def test_transpose_op_with_axes(op_class):\n", + " x = pt.tensor(\"x\", shape=(2, None, 4))\n", + " x_test = np.arange(2 * 3 * 4).reshape((2, 3, 4))\n", + "\n", + " for axis, dtype in [\n", + " ((0, 2, 1), \"int64\"),\n", + " ((2, 0, 1), \"float32\")]:\n", + " op = op_class(axis)\n", + " out = op(x.astype(dtype))\n", + "\n", + " assert out.type.ndim == 3\n", + " assert out.type.dtype == dtype\n", + " np.testing.assert_allclose(out.eval({x: x_test}), x_test.transpose(axis))\n", + "\n", + "# test_transpose_op_with_axes(Transpose) # uncomment me" + ] + }, + { + "cell_type": "markdown", + "id": "35c78a38-4ce7-4514-beff-4d926999beab", + "metadata": { + "id": "35c78a38-4ce7-4514-beff-4d926999beab" + }, + "source": [ + "### Exercise 3: Define operator equality using `__props__`\n", + "\n", + "PyTensor tries to avoid recomputing equivalent computations in a graph. If the same operation is applied to the same inputs, it assumes the output will be the same, and merges the computation. Here is an example using the Sum axis" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "4d6c15d1-976e-4387-b1dc-dc7dc031d315", + "metadata": { + "id": "4d6c15d1-976e-4387-b1dc-dc7dc031d315" + }, + "outputs": [], + "source": [ + "x = pt.vector(\"x\")\n", + "out = sum(x) + sum(x)" + ] + }, + { + "cell_type": "markdown", + "id": "5c26ad47-ab04-40f7-8715-33a5d5ce5417", + "metadata": { + "id": "5c26ad47-ab04-40f7-8715-33a5d5ce5417" + }, + "source": [ + "The original graph contains 2 distinct Sum operations (note the different ids)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "76fd85ba-d57d-4195-a030-00e0b0bf6bf2", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "76fd85ba-d57d-4195-a030-00e0b0bf6bf2", + "outputId": "b7d0d831-e88f-4cbd-f526-555a5bddef9a" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Add [id A]\n", + " โ”œโ”€ Sum [id B]\n", + " โ”‚ โ””โ”€ x [id C]\n", + " โ””โ”€ Sum [id D]\n", + " โ””โ”€ x [id C]\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "out.dprint()" + ] + }, + { + "cell_type": "markdown", + "id": "2173f4ee-d2a1-4aea-b01d-7f452f7df79b", + "metadata": { + "id": "2173f4ee-d2a1-4aea-b01d-7f452f7df79b" + }, + "source": [ + "But after rewriting only one sum is computed (note the same ids and the ellipsis)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "30cd0072-1ecf-428e-82ca-4a094eb3bd1a", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "30cd0072-1ecf-428e-82ca-4a094eb3bd1a", + "outputId": "52a6056a-5620-4763-a0f6-73f0995b3321" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Add [id A]\n", + " โ”œโ”€ Sum [id B]\n", + " โ”‚ โ””โ”€ x [id C]\n", + " โ””โ”€ Sum [id B]\n", + " โ””โ”€ ยทยทยท\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "rewrite_graph(out).dprint()" + ] + }, + { + "cell_type": "markdown", + "id": "bb99df31-120b-43b0-8371-b101c5c74011", + "metadata": { + "id": "bb99df31-120b-43b0-8371-b101c5c74011" + }, + "source": [ + "However if we use different instances of the Sum Op PyTensor does not consider them equivalent and no merging is done." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "53b475a6-bdd1-4783-8cd0-d08b48d78adc", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "53b475a6-bdd1-4783-8cd0-d08b48d78adc", + "outputId": "de25156a-40c1-4775-85ac-d8b3ae01519f" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Add [id A]\n", + " โ”œโ”€ Sum [id B]\n", + " โ”‚ โ””โ”€ x [id C]\n", + " โ””โ”€ Sum [id D]\n", + " โ””โ”€ x [id C]\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "out = Sum()(x) + Sum()(x)\n", + "rewrite_graph(out).dprint()" + ] + }, + { + "cell_type": "markdown", + "id": "058eb9a1-ee07-4c2a-82db-1b5074340366", + "metadata": { + "id": "058eb9a1-ee07-4c2a-82db-1b5074340366" + }, + "source": [ + "PyTensor uses Op equality to determine if two computations are equivalent. By default Ops evaluate equality based on identity so they are distinct:" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "de6d2242-d7ef-423c-b98c-e52cb3a60b19", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, + "id": "de6d2242-d7ef-423c-b98c-e52cb3a60b19", + "outputId": "861543c6-2428-4307-825e-93901af8aa7c" + }, + "outputs": [ { - "cell_type": "code", - "execution_count": 20, - "id": "aeee3620-ea84-4027-a5c3-c1d20fecdb08", - "metadata": { - "id": "aeee3620-ea84-4027-a5c3-c1d20fecdb08" - }, - "outputs": [], - "source": [ - "class Convolve(Op):\n", - " ...\n", - "\n", - "def test_convolve(op_class):\n", - " x = pt.vector(\"x\", shape=(None,))\n", - " y = pt.vector(\"y\", shape=(3,))\n", - " out = op_class()(x, y)\n", - "\n", - " x_test = np.arange(10).astype(\"float64\")\n", - " y_test = np.array([0, 1, 2]).astype=(\"float64\")\n", - " res = out.eval({x: x_test, y: y_test})\n", - "\n", - " np.testing.assert_allclose(res, np.convolve(x_test, y_test))\n", - "\n", - " res2 = out.eval({x: res, y: y_test})\n", - " np.testing.assert_allclose(res, np.convolve(res, y_test))\n", - "\n", - "# test_convolve(Convolve) # uncomment me" + "data": { + "text/plain": [ + "False" ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "Sum() == Sum()" + ] + }, + { + "cell_type": "markdown", + "id": "03f5da98-26cd-4753-9809-a50fc89e58c7", + "metadata": { + "id": "03f5da98-26cd-4753-9809-a50fc89e58c7" + }, + "source": [ + "This is not the case for the PyTensor implementation of Sum" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "7b5963fd-ae84-40a7-9fbe-8a65a76fef51", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, + "id": "7b5963fd-ae84-40a7-9fbe-8a65a76fef51", + "outputId": "ddab8167-3d51-49ca-abc1-d922401b261a" + }, + "outputs": [ { - "cell_type": "markdown", - "id": "46cc34d2-fa6b-448e-a810-55158a85485d", - "metadata": { - "id": "46cc34d2-fa6b-448e-a810-55158a85485d" - }, - "source": [ - "Extend the Op to include the parameter `mode` that `np.convolve` also offers.\n", - "\n", - "Extra points if the output shape is specified when that's possible" + "data": { + "text/plain": [ + "True" ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pt.sum(x).owner.op == pt.sum(x).owner.op" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "e26aa404-0c89-4ee0-a370-f080df6c9584", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "e26aa404-0c89-4ee0-a370-f080df6c9584", + "outputId": "a0badf04-7918-4a30-b46a-5a2a99efaa8f" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Add [id A]\n", + " โ”œโ”€ Sum{axes=None} [id B]\n", + " โ”‚ โ””โ”€ x [id C]\n", + " โ””โ”€ Sum{axes=None} [id B]\n", + " โ””โ”€ ยทยทยท\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "rewrite_graph(pt.sum(x) + pt.sum(x)).dprint()" + ] + }, + { + "cell_type": "markdown", + "id": "2e6d2797-b70d-4ef8-83c1-9947014654bf", + "metadata": { + "id": "2e6d2797-b70d-4ef8-83c1-9947014654bf" + }, + "source": [ + "The default way of implementing Op equality is to define `__props__`, a tuple of strings with the names of immutable instance properties that \"parametrize\" an `Op`.\n", + "\n", + "When an `Op` has `__props__`, PyTensor will check if the respective instance attributes are equal and if so, assume two Operations from the same class are equivalent.\n", + "\n", + "Our simplest implementation of Sum has no parametrization, so we can define an empty `__props__`:" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "99ba7e22-17dd-4ec6-aaca-d281699877df", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, + "id": "99ba7e22-17dd-4ec6-aaca-d281699877df", + "outputId": "8c5bcca7-ebf4-4e9a-ffed-aa1d483902c9" + }, + "outputs": [ { - "cell_type": "code", - "execution_count": 21, - "id": "355acf46-d232-4368-9d96-e63378da3d9d", - "metadata": { - "id": "355acf46-d232-4368-9d96-e63378da3d9d" - }, - "outputs": [], - "source": [ - "class Convolve(Op):\n", - " ...\n", - "\n", - "def test_convolve(op_class):\n", - " x = pt.vector(\"x\", shape=(10,))\n", - " y = pt.vector(\"y\", shape=(3,))\n", - "\n", - " x_test = np.arange(10).astype(\"float64\")\n", - " y_test = np.array([0, 1, 2]).astype=(\"float64\")\n", - "\n", - " for mode in (\"full\", \"valid\", \"same\"):\n", - " print(f\"{mode=}\")\n", - " op = op_class(mode=mode)\n", - " assert op == op_class(mode=mode)\n", - "\n", - " out = op(x, y)\n", - " if out.type.shape != (None,):\n", - " assert out.type.shape == np.convolve(x_test, y_test, mode=mode).shape\n", - "\n", - "\n", - " res = out.eval({x: x_test, y: y_test})\n", - " np.testing.assert_allclose(res, np.convolve(x_test, y_test, mode=mode))\n", - "\n", - "# test_convolve(Convolve) # uncomment me" + "data": { + "text/plain": [ + "True" ] - }, - { - "cell_type": "markdown", - "source": [ - "Open-ended challenge: implement an Op of your choosing.\n", - "\n", - "Some ideas of Ops that don't currently exist in PyTensor:\n", - "* [numpy.frexp](https://numpy.org/doc/2.1/reference/generated/numpy.frexp.html)\n", - "* [numpy.nextafter](https://numpy.org/doc/2.1/reference/generated/numpy.nextafter.html)\n", - "* [scipy.special.gauss_spline](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.gauss_spline.html#scipy.signal.gauss_spline)\n", - "* Anything else you fancy" - ], - "metadata": { - "id": "VVqIMwqZ2xIC" - }, - "id": "VVqIMwqZ2xIC" - }, - { - "cell_type": "code", - "source": [], - "metadata": { - "id": "NVa-_DtL4sOC" - }, - "id": "NVa-_DtL4sOC", - "execution_count": 21, - "outputs": [] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" } - ], - "metadata": { - "kernelspec": { - "display_name": "pytensor", - "language": "python", - "name": "pytensor" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.8" - }, + ], + "source": [ + "class Sum(Op):\n", + " __props__ = ()\n", + "\n", + " def make_node(self, x):\n", + " return Apply(self, [x], [pt.scalar()])\n", + "\n", + " def perform(self, node, inputs, outputs):\n", + " outputs[0][0] = inputs[0].sum()\n", + "\n", + "Sum() == Sum()" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "8f09f978-1077-4a4c-97f9-f0a7df4dfe61", + "metadata": { "colab": { - "provenance": [], - "toc_visible": true, - "include_colab_link": true + "base_uri": "https://localhost:8080/" + }, + "id": "8f09f978-1077-4a4c-97f9-f0a7df4dfe61", + "outputId": "a1c27444-7513-4eb9-94e2-6e6b626e9344" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Add [id A]\n", + " โ”œโ”€ Sum [id B]\n", + " โ”‚ โ””โ”€ x [id C]\n", + " โ””โ”€ Sum [id B]\n", + " โ””โ”€ ยทยทยท\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" } + ], + "source": [ + "rewrite_graph(Sum()(x) + Sum()(x)).dprint()" + ] + }, + { + "cell_type": "markdown", + "id": "81a849a3-ca89-43d7-b8ba-c47420e56853", + "metadata": { + "id": "81a849a3-ca89-43d7-b8ba-c47420e56853" + }, + "source": [ + "Extend the Transpose Op with `__props__` so that two instances with the same axis evaluate equal." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "1259f680-a458-4c6b-b96e-9fabdf2a938e", + "metadata": { + "id": "1259f680-a458-4c6b-b96e-9fabdf2a938e" + }, + "outputs": [], + "source": [ + "class Transpose(Op):\n", + " ...\n", + "\n", + "@test\n", + "def test_transpose_op_with_axes_and_props(op_class):\n", + " x = pt.tensor(\"x\", shape=(2, None, 4))\n", + " x_test = np.arange(2 * 3 * 4).reshape((2, 3, 4))\n", + "\n", + " assert len(op_class.__props__)\n", + " assert op_class(axis=(0, 2, 1)) == op_class(axis=(0, 2, 1))\n", + " assert op_class(axis=(0, 2, 1)) != op_class(axis=(2, 0, 1))\n", + "\n", + "# test_transpose_op_with_axes_and_props(Transpose) # uncomment me" + ] + }, + { + "cell_type": "markdown", + "id": "000c0f3b-18bd-4494-a4a5-9010a03cade5", + "metadata": { + "id": "000c0f3b-18bd-4494-a4a5-9010a03cade5" + }, + "source": [ + "### Exercise 4, implement an Op that wraps `np.convolve`" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "aeee3620-ea84-4027-a5c3-c1d20fecdb08", + "metadata": { + "id": "aeee3620-ea84-4027-a5c3-c1d20fecdb08" + }, + "outputs": [], + "source": [ + "class Convolve(Op):\n", + " ...\n", + "\n", + "def test_convolve(op_class):\n", + " x = pt.vector(\"x\", shape=(None,))\n", + " y = pt.vector(\"y\", shape=(3,))\n", + " out = op_class()(x, y)\n", + "\n", + " x_test = np.arange(10).astype(\"float64\")\n", + " y_test = np.array([0, 1, 2]).astype=(\"float64\")\n", + " res = out.eval({x: x_test, y: y_test})\n", + "\n", + " np.testing.assert_allclose(res, np.convolve(x_test, y_test))\n", + "\n", + " res2 = out.eval({x: res, y: y_test})\n", + " np.testing.assert_allclose(res, np.convolve(res, y_test))\n", + "\n", + "# test_convolve(Convolve) # uncomment me" + ] + }, + { + "cell_type": "markdown", + "id": "46cc34d2-fa6b-448e-a810-55158a85485d", + "metadata": { + "id": "46cc34d2-fa6b-448e-a810-55158a85485d" + }, + "source": [ + "Extend the Op to include the parameter `mode` that `np.convolve` also offers.\n", + "\n", + "Extra points if the output shape is specified when that's possible" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "355acf46-d232-4368-9d96-e63378da3d9d", + "metadata": { + "id": "355acf46-d232-4368-9d96-e63378da3d9d" + }, + "outputs": [], + "source": [ + "class Convolve(Op):\n", + " ...\n", + "\n", + "def test_convolve(op_class):\n", + " x = pt.vector(\"x\", shape=(10,))\n", + " y = pt.vector(\"y\", shape=(3,))\n", + "\n", + " x_test = np.arange(10).astype(\"float64\")\n", + " y_test = np.array([0, 1, 2]).astype=(\"float64\")\n", + "\n", + " for mode in (\"full\", \"valid\", \"same\"):\n", + " print(f\"{mode=}\")\n", + " op = op_class(mode=mode)\n", + " assert op == op_class(mode=mode)\n", + "\n", + " out = op(x, y)\n", + " if out.type.shape != (None,):\n", + " assert out.type.shape == np.convolve(x_test, y_test, mode=mode).shape\n", + "\n", + "\n", + " res = out.eval({x: x_test, y: y_test})\n", + " np.testing.assert_allclose(res, np.convolve(x_test, y_test, mode=mode))\n", + "\n", + "# test_convolve(Convolve) # uncomment me" + ] + }, + { + "cell_type": "markdown", + "id": "VVqIMwqZ2xIC", + "metadata": { + "id": "VVqIMwqZ2xIC" + }, + "source": [ + "Open-ended challenge: implement an Op of your choosing.\n", + "\n", + "Some ideas of Ops that don't currently exist in PyTensor:\n", + "* [numpy.frexp](https://numpy.org/doc/2.1/reference/generated/numpy.frexp.html)\n", + "* [numpy.nextafter](https://numpy.org/doc/2.1/reference/generated/numpy.nextafter.html)\n", + "* [scipy.special.gauss_spline](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.gauss_spline.html#scipy.signal.gauss_spline)\n", + "* Anything else you fancy" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "NVa-_DtL4sOC", + "metadata": { + "id": "NVa-_DtL4sOC" + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "colab": { + "include_colab_link": true, + "provenance": [], + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.14" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/scripts/pytensor_from_scratch.py b/scripts/pytensor_from_scratch.py new file mode 100644 index 0000000..7bdda98 --- /dev/null +++ b/scripts/pytensor_from_scratch.py @@ -0,0 +1,100 @@ +""" +This script is an export of pytensor_from_scratch.ipynb. +It was created to generate the UML and call charts in data/. +""" + + +class Type: + "Baseclass for PyTensor types" + +class Op: + "Baseclass for PyTensor operations." + + def __str__(self): + return self.__class__.__name__ + +class Node: + "Baseclass for PyTensor nodes." + + +class Variable(Node): + def __init__(self, name=None, *, type: Type): + self.name = name + self.type = type + self.owner = None + + def __repr__(self): + if self.name: + return self.name + return f"Variable(type={self.type})" + +class Apply(Node): + def __init__(self, op:Op, inputs, outputs): + self.op = op + self.inputs = inputs + self.outputs = outputs + for out in outputs: + if out.owner is not None: + raise ValueError("This variable already belongs to another Apply Node") + out.owner = self + + def __repr__(self): + return f"Apply(op={self.op.__class__.__name__}, inputs={self.inputs}, outputs={self.outputs})" + + +class TensorType(Type): + def __init__(self, shape: tuple[float | None, ...], dtype: str): + self.shape = shape + self.dtype = dtype + + def __eq__(self, other): + return ( + type(self) is type(other) + and self.shape == other.shape + and self.dtype == other.dtype + ) + + def __repr__(self): + return f"TensorType(shape={self.shape}, dtype={self.dtype})" + +class Add(Op): + def make_node(self, a, b): + if not(isinstance(a.type, TensorType) and isinstance(b.type, TensorType)): + raise TypeError("Inputs must be tensors") + if a.type != b.type: + raise TypeError("Addition only supported for inputs of the same type") + + output_type = TensorType(shape=a.type.shape, dtype=a.type.dtype) + output = Variable(type=output_type) + return Apply(self, [a, b], [output]) + + +add = Add() + +dvector = TensorType(shape=(10,), dtype="float64") + +x = Variable("x", type=dvector) +y = Variable("y", type=dvector) +[x_add_y] = add.make_node(x, y).outputs +x_add_y.name = "x + y" +x_add_y.owner + +class Sum(Op): + + def make_node(self, a): + if not(isinstance(a.type, TensorType)): + raise TypeError("Input must be a tensor") + output_type = TensorType(shape=(), dtype=a.type.dtype) + output = Variable(type=output_type) + return Apply(self, [a], [output]) + +sum = Sum() + +[sum_x_add_y] = sum.make_node(x_add_y).outputs +sum_x_add_y.name = "sum(x + y)" +sum_x_add_y.owner + + + + +