numba_inplace_bug.ipynb

{
  "cells" : [ {
    "metadata" : {
      "ExecuteTime" : {
        "end_time" : "2025-02-23T07:46:13.431226Z",
        "start_time" : "2025-02-23T07:46:13.428335Z"
      }
    },
    "cell_type" : "code",
    "source" : [ "import numpy as np\n", "import pytensor\n", "import pytensor.tensor as pt\n", "from pytensor.compile.mode import get_mode" ],
    "id" : "d9aab46b5c33cf3",
    "outputs" : [ ],
    "execution_count" : 13
  }, {
    "metadata" : {
      "collapsed" : true,
      "ExecuteTime" : {
        "end_time" : "2025-02-23T07:43:11.212010Z",
        "start_time" : "2025-02-23T07:43:10.699012Z"
      }
    },
    "cell_type" : "code",
    "source" : [ "rng = np.random.default_rng(3)\n", "l_test = rng.normal(size=(5, 5))\n", "x_test = l_test @ l_test.T\n", "\n", "x = pt.tensor(\"x\", shape=x_test.shape)\n", "x_chol = pt.linalg.cholesky(x, lower=True)\n", "x_chol_grad = grad = pt.grad(x_chol.sum(), x)\n", "\n", "fn = pytensor.function([x], x_chol_grad)\n", "fn.dprint(print_type=True, print_destroy_map=True, print_view_map=True)" ],
    "id" : "initial_id",
    "outputs" : [ {
      "name" : "stdout",
      "output_type" : "stream",
      "text" : [ "Composite{(((i0 + i1) * i2) - i3)} [id A] <Matrix(float64, shape=(5, 5))> 15\n", " ├─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id B] <Matrix(float64, shape=(5, 5))> d={0: [1]} 11\n", " │  ├─ Transpose{axes=[1, 0]} [id C] <Matrix(float64, shape=(5, 5))> v={0: [0]} 1\n", " │  │  └─ Cholesky{lower=True, check_finite=True, on_error='raise', overwrite_a=False} [id D] <Matrix(float64, shape=(5, 5))> 0\n", " │  │     └─ x [id E] <Matrix(float64, shape=(5, 5))>\n", " │  └─ Transpose{axes=[1, 0]} [id F] <Matrix(float64, shape=(5, 5))> v={0: [0]} 10\n", " │     └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id G] <Matrix(float64, shape=(5, 5))> d={0: [1]} 9\n", " │        ├─ Transpose{axes=[1, 0]} [id C] <Matrix(float64, shape=(5, 5))> v={0: [0]} 1\n", " │        │  └─ ···\n", " │        └─ Composite{((i0 * i1) - i2)} [id H] <Matrix(float64, shape=(5, 5))> d={0: [0]} 8\n", " │           ├─ Dot22 [id I] <Matrix(float64, shape=(5, 5))> 2\n", " │           │  ├─ [[1. 1. 1. ... 1. 1. 1.]] [id J] <Matrix(float64, shape=(5, 5))>\n", " │           │  └─ Cholesky{lower=True, check_finite=True, on_error='raise', overwrite_a=False} [id D] <Matrix(float64, shape=(5, 5))> 0\n", " │           │     └─ ···\n", " │           ├─ [[1. 1. 1. ... 0. 0. 1.]] [id K] <Matrix(float64, shape=(5, 5))>\n", " │           └─ Transpose{axes=[1, 0]} [id L] <Matrix(float64, shape=(?, ?))> v={0: [0]} 7\n", " │              └─ AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id M] <Matrix(float64, shape=(?, ?))> 6\n", " │                 └─ Mul [id N] <Vector(float64, shape=(?,))> d={0: [1]} 5\n", " │                    ├─ [0.5] [id O] <Vector(float64, shape=(1,))>\n", " │                    └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id P] <Vector(float64, shape=(?,))> 4\n", " │                       └─ Dot22 [id Q] <Matrix(float64, shape=(5, 5))> 3\n", " │                          ├─ Transpose{axes=[1, 0]} [id C] <Matrix(float64, shape=(5, 5))> v={0: [0]} 1\n", " │                          │  └─ ···\n", " │                          └─ [[1. 1. 1. ... 1. 1. 1.]] [id J] <Matrix(float64, shape=(5, 5))>\n", " ├─ Transpose{axes=[1, 0]} [id R] <Matrix(float64, shape=(5, 5))> v={0: [0]} 13\n", " │  └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id B] <Matrix(float64, shape=(5, 5))> d={0: [1]} 11\n", " │     └─ ···\n", " ├─ [[1. 0. 0. ... 1. 1. 1.]] [id S] <Matrix(float64, shape=(5, 5))>\n", " └─ AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id T] <Matrix(float64, shape=(?, ?))> 14\n", "    └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id U] <Vector(float64, shape=(?,))> 12\n", "       └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id B] <Matrix(float64, shape=(5, 5))> d={0: [1]} 11\n", "          └─ ···\n", "\n", "Inner graphs:\n", "\n", "Composite{(((i0 + i1) * i2) - i3)} [id A]\n", " ← sub [id V] <float64> 'o0'\n", "    ├─ mul [id W] <float64>\n", "    │  ├─ add [id X] <float64>\n", "    │  │  ├─ i0 [id Y] <float64>\n", "    │  │  └─ i1 [id Z] <float64>\n", "    │  └─ i2 [id BA] <float64>\n", "    └─ i3 [id BB] <float64>\n", "\n", "Composite{((i0 * i1) - i2)} [id H] d={0: [0]}\n", " ← sub [id BC] <float64> 'o0'\n", "    ├─ mul [id BD] <float64>\n", "    │  ├─ i0 [id BE] <float64>\n", "    │  └─ i1 [id BF] <float64>\n", "    └─ i2 [id BG] <float64>\n", "\n", "AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id M]\n", " ← AdvancedSetSubtensor [id BH] <Matrix(float64, shape=(?, ?))>\n", "    ├─ Alloc [id BI] <Matrix(float64, shape=(?, ?))>\n", "    │  ├─ 0.0 [id BJ] <Scalar(float64, shape=())>\n", "    │  ├─ Add [id BK] <Scalar(int64, shape=())>\n", "    │  │  ├─ Subtensor{i} [id BL] <Scalar(int64, shape=())> v={0: [0]}\n", "    │  │  │  ├─ Shape [id BM] <Vector(int64, shape=(1,))>\n", "    │  │  │  │  └─ *0-<Vector(float64, shape=(?,))> [id BN] <Vector(float64, shape=(?,))>\n", "    │  │  │  └─ -1 [id BO] <int64>\n", "    │  │  └─ 0 [id BP] <Scalar(int8, shape=())>\n", "    │  └─ Add [id BK] <Scalar(int64, shape=())>\n", "    │     └─ ···\n", "    ├─ *0-<Vector(float64, shape=(?,))> [id BN] <Vector(float64, shape=(?,))>\n", "    ├─ Add [id BQ] <Vector(int64, shape=(?,))>\n", "    │  ├─ ARange{dtype='int64'} [id BR] <Vector(int64, shape=(?,))>\n", "    │  │  ├─ 0 [id BS] <Scalar(int8, shape=())>\n", "    │  │  ├─ Subtensor{i} [id BT] <Scalar(int64, shape=())> v={0: [0]}\n", "    │  │  │  ├─ Shape [id BU] <Vector(int64, shape=(1,))>\n", "    │  │  │  │  └─ *0-<Vector(float64, shape=(?,))> [id BN] <Vector(float64, shape=(?,))>\n", "    │  │  │  └─ -1 [id BV] <int64>\n", "    │  │  └─ 1 [id BW] <Scalar(int8, shape=())>\n", "    │  └─ ExpandDims{axis=0} [id BX] <Vector(int64, shape=(1,))> v={0: [0]}\n", "    │     └─ 0 [id BY] <Scalar(int64, shape=())>\n", "    └─ Add [id BZ] <Vector(int64, shape=(?,))>\n", "       ├─ ARange{dtype='int64'} [id BR] <Vector(int64, shape=(?,))>\n", "       │  └─ ···\n", "       └─ ExpandDims{axis=0} [id CA] <Vector(int64, shape=(1,))> v={0: [0]}\n", "          └─ 0 [id CB] <Scalar(int64, shape=())>\n", "\n", "AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id T]\n", " ← AdvancedSetSubtensor [id CC] <Matrix(float64, shape=(?, ?))>\n", "    ├─ Alloc [id CD] <Matrix(float64, shape=(?, ?))>\n", "    │  ├─ 0.0 [id CE] <Scalar(float64, shape=())>\n", "    │  ├─ Add [id CF] <Scalar(int64, shape=())>\n", "    │  │  ├─ Subtensor{i} [id CG] <Scalar(int64, shape=())> v={0: [0]}\n", "    │  │  │  ├─ Shape [id CH] <Vector(int64, shape=(1,))>\n", "    │  │  │  │  └─ *0-<Vector(float64, shape=(?,))> [id BN] <Vector(float64, shape=(?,))>\n", "    │  │  │  └─ -1 [id CI] <int64>\n", "    │  │  └─ 0 [id CJ] <Scalar(int8, shape=())>\n", "    │  └─ Add [id CF] <Scalar(int64, shape=())>\n", "    │     └─ ···\n", "    ├─ *0-<Vector(float64, shape=(?,))> [id BN] <Vector(float64, shape=(?,))>\n", "    ├─ Add [id CK] <Vector(int64, shape=(?,))>\n", "    │  ├─ ARange{dtype='int64'} [id CL] <Vector(int64, shape=(?,))>\n", "    │  │  ├─ 0 [id CM] <Scalar(int8, shape=())>\n", "    │  │  ├─ Subtensor{i} [id CN] <Scalar(int64, shape=())> v={0: [0]}\n", "    │  │  │  ├─ Shape [id CO] <Vector(int64, shape=(1,))>\n", "    │  │  │  │  └─ *0-<Vector(float64, shape=(?,))> [id BN] <Vector(float64, shape=(?,))>\n", "    │  │  │  └─ -1 [id CP] <int64>\n", "    │  │  └─ 1 [id CQ] <Scalar(int8, shape=())>\n", "    │  └─ ExpandDims{axis=0} [id CR] <Vector(int64, shape=(1,))> v={0: [0]}\n", "    │     └─ 0 [id CS] <Scalar(int64, shape=())>\n", "    └─ Add [id CT] <Vector(int64, shape=(?,))>\n", "       ├─ ARange{dtype='int64'} [id CL] <Vector(int64, shape=(?,))>\n", "       │  └─ ···\n", "       └─ ExpandDims{axis=0} [id CU] <Vector(int64, shape=(1,))> v={0: [0]}\n", "          └─ 0 [id CV] <Scalar(int64, shape=())>\n" ]
    }, {
      "data" : {
        "text/plain" : [ "<ipykernel.iostream.OutStream at 0x7fc6f8a78a60>" ]
      },
      "execution_count" : 8,
      "metadata" : { },
      "output_type" : "execute_result"
    } ],
    "execution_count" : 8
  }, {
    "metadata" : {
      "ExecuteTime" : {
        "end_time" : "2025-02-23T07:43:22.583215Z",
        "start_time" : "2025-02-23T07:43:22.579455Z"
      }
    },
    "cell_type" : "code",
    "source" : "fn(x_test).round(2)",
    "id" : "98c2729d4d9135fe",
    "outputs" : [ {
      "data" : {
        "text/plain" : [ "array([[ 0.42,  0.  ,  0.  ,  0.  ,  0.  ],\n", "       [ 1.24,  1.05,  0.  ,  0.  ,  0.  ],\n", "       [ 2.19,  3.69,  3.39,  0.  ,  0.  ],\n", "       [ 3.72,  6.08, 11.28, 10.09,  0.  ],\n", "       [ 1.7 ,  3.15,  6.06,  9.87,  2.21]])" ]
      },
      "execution_count" : 10,
      "metadata" : { },
      "output_type" : "execute_result"
    } ],
    "execution_count" : 10
  }, {
    "metadata" : {
      "ExecuteTime" : {
        "end_time" : "2025-02-23T07:47:07.596976Z",
        "start_time" : "2025-02-23T07:47:07.374815Z"
      }
    },
    "cell_type" : "code",
    "source" : [ "fn = pytensor.function([x], x_chol_grad, mode=get_mode(\"NUMBA\").including(\"inplace\"))\n", "fn.dprint(print_type=True, print_destroy_map=True, print_view_map=True)" ],
    "id" : "9834bb6d8d5a7e91",
    "outputs" : [ {
      "name" : "stdout",
      "output_type" : "stream",
      "text" : [ "Composite{(((i0 + i1) * i2) - i3)} [id A] <Matrix(float64, shape=(5, 5))> 15\n", " ├─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id B] <Matrix(float64, shape=(5, 5))> d={0: [1]} 11\n", " │  ├─ Transpose{axes=[1, 0]} [id C] <Matrix(float64, shape=(5, 5))> v={0: [0]} 1\n", " │  │  └─ Cholesky{lower=True, check_finite=True, on_error='raise', overwrite_a=False} [id D] <Matrix(float64, shape=(5, 5))> 0\n", " │  │     └─ x [id E] <Matrix(float64, shape=(5, 5))>\n", " │  └─ Transpose{axes=[1, 0]} [id F] <Matrix(float64, shape=(5, 5))> v={0: [0]} 10\n", " │     └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id G] <Matrix(float64, shape=(5, 5))> d={0: [1]} 9\n", " │        ├─ Transpose{axes=[1, 0]} [id C] <Matrix(float64, shape=(5, 5))> v={0: [0]} 1\n", " │        │  └─ ···\n", " │        └─ Composite{((i0 * i1) - i2)} [id H] <Matrix(float64, shape=(5, 5))> d={0: [0]} 8\n", " │           ├─ dot [id I] <Matrix(float64, shape=(5, 5))> 2\n", " │           │  ├─ [[1. 1. 1. ... 1. 1. 1.]] [id J] <Matrix(float64, shape=(5, 5))>\n", " │           │  └─ Cholesky{lower=True, check_finite=True, on_error='raise', overwrite_a=False} [id D] <Matrix(float64, shape=(5, 5))> 0\n", " │           │     └─ ···\n", " │           ├─ [[1. 1. 1. ... 0. 0. 1.]] [id K] <Matrix(float64, shape=(5, 5))>\n", " │           └─ Transpose{axes=[1, 0]} [id L] <Matrix(float64, shape=(?, ?))> v={0: [0]} 7\n", " │              └─ AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id M] <Matrix(float64, shape=(?, ?))> 6\n", " │                 └─ Mul [id N] <Vector(float64, shape=(?,))> d={0: [1]} 5\n", " │                    ├─ [0.5] [id O] <Vector(float64, shape=(1,))>\n", " │                    └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id P] <Vector(float64, shape=(?,))> 4\n", " │                       └─ dot [id Q] <Matrix(float64, shape=(5, 5))> 3\n", " │                          ├─ Transpose{axes=[1, 0]} [id C] <Matrix(float64, shape=(5, 5))> v={0: [0]} 1\n", " │                          │  └─ ···\n", " │                          └─ [[1. 1. 1. ... 1. 1. 1.]] [id J] <Matrix(float64, shape=(5, 5))>\n", " ├─ Transpose{axes=[1, 0]} [id R] <Matrix(float64, shape=(5, 5))> v={0: [0]} 13\n", " │  └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id B] <Matrix(float64, shape=(5, 5))> d={0: [1]} 11\n", " │     └─ ···\n", " ├─ [[1. 0. 0. ... 1. 1. 1.]] [id S] <Matrix(float64, shape=(5, 5))>\n", " └─ AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id T] <Matrix(float64, shape=(?, ?))> 14\n", "    └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id U] <Vector(float64, shape=(?,))> 12\n", "       └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True} [id B] <Matrix(float64, shape=(5, 5))> d={0: [1]} 11\n", "          └─ ···\n", "\n", "Inner graphs:\n", "\n", "Composite{(((i0 + i1) * i2) - i3)} [id A]\n", " ← sub [id V] <float64> 'o0'\n", "    ├─ mul [id W] <float64>\n", "    │  ├─ add [id X] <float64>\n", "    │  │  ├─ i0 [id Y] <float64>\n", "    │  │  └─ i1 [id Z] <float64>\n", "    │  └─ i2 [id BA] <float64>\n", "    └─ i3 [id BB] <float64>\n", "\n", "Composite{((i0 * i1) - i2)} [id H] d={0: [0]}\n", " ← sub [id BC] <float64> 'o0'\n", "    ├─ mul [id BD] <float64>\n", "    │  ├─ i0 [id BE] <float64>\n", "    │  └─ i1 [id BF] <float64>\n", "    └─ i2 [id BG] <float64>\n", "\n", "AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id M]\n", " ← AdvancedSetSubtensor [id BH] <Matrix(float64, shape=(?, ?))> d={0: [0]}\n", "    ├─ Alloc [id BI] <Matrix(float64, shape=(?, ?))>\n", "    │  ├─ 0.0 [id BJ] <Scalar(float64, shape=())>\n", "    │  ├─ Shape_i{0} [id BK] <Scalar(int64, shape=())>\n", "    │  │  └─ *0-<Vector(float64, shape=(?,))> [id BL] <Vector(float64, shape=(?,))>\n", "    │  └─ Shape_i{0} [id BK] <Scalar(int64, shape=())>\n", "    │     └─ ···\n", "    ├─ *0-<Vector(float64, shape=(?,))> [id BL] <Vector(float64, shape=(?,))>\n", "    ├─ ARange{dtype='int64'} [id BM] <Vector(int64, shape=(?,))>\n", "    │  ├─ 0 [id BN] <Scalar(int8, shape=())>\n", "    │  ├─ Shape_i{0} [id BK] <Scalar(int64, shape=())>\n", "    │  │  └─ ···\n", "    │  └─ 1 [id BO] <Scalar(int8, shape=())>\n", "    └─ ARange{dtype='int64'} [id BM] <Vector(int64, shape=(?,))>\n", "       └─ ···\n", "\n", "AllocDiag{self.axis1=0, self.axis2=1, self.offset=0} [id T]\n", " ← AdvancedSetSubtensor [id BP] <Matrix(float64, shape=(?, ?))> d={0: [0]}\n", "    ├─ Alloc [id BQ] <Matrix(float64, shape=(?, ?))>\n", "    │  ├─ 0.0 [id BR] <Scalar(float64, shape=())>\n", "    │  ├─ Shape_i{0} [id BS] <Scalar(int64, shape=())>\n", "    │  │  └─ *0-<Vector(float64, shape=(?,))> [id BL] <Vector(float64, shape=(?,))>\n", "    │  └─ Shape_i{0} [id BS] <Scalar(int64, shape=())>\n", "    │     └─ ···\n", "    ├─ *0-<Vector(float64, shape=(?,))> [id BL] <Vector(float64, shape=(?,))>\n", "    ├─ ARange{dtype='int64'} [id BT] <Vector(int64, shape=(?,))>\n", "    │  ├─ 0 [id BU] <Scalar(int8, shape=())>\n", "    │  ├─ Shape_i{0} [id BS] <Scalar(int64, shape=())>\n", "    │  │  └─ ···\n", "    │  └─ 1 [id BV] <Scalar(int8, shape=())>\n", "    └─ ARange{dtype='int64'} [id BT] <Vector(int64, shape=(?,))>\n", "       └─ ···\n" ]
    }, {
      "name" : "stderr",
      "output_type" : "stream",
      "text" : [ "/home/ricardo/Documents/pytensor/pytensor/graph/rewriting/basic.py:120: UserWarning: A Supervisor feature is missing from FunctionGraph(AdvancedSetSubtensor(Alloc(0.0, *2 -> Shape_i{0}(*0-<Vector(float64, shape=(?,))>), *2), *0-<Vector(float64, shape=(?,))>, *1 -> ARange{dtype='int64'}(0, *2, 1), *1)).\n", "  return self.apply(fgraph, *args, **kwargs)\n", "/home/ricardo/Documents/pytensor/pytensor/graph/rewriting/basic.py:120: UserWarning: A Supervisor feature is missing from FunctionGraph(AdvancedSetSubtensor(Alloc(0.0, *2 -> Shape_i{0}(*0-<Vector(float64, shape=(?,))>), *2), *0-<Vector(float64, shape=(?,))>, *1 -> ARange{dtype='int64'}(0, *2, 1), *1)).\n", "  return self.apply(fgraph, *args, **kwargs)\n" ]
    }, {
      "data" : {
        "text/plain" : [ "<ipykernel.iostream.OutStream at 0x7fc6f8a78a60>" ]
      },
      "execution_count" : 16,
      "metadata" : { },
      "output_type" : "execute_result"
    } ],
    "execution_count" : 16
  }, {
    "metadata" : {
      "ExecuteTime" : {
        "end_time" : "2025-02-23T07:47:18.232244Z",
        "start_time" : "2025-02-23T07:47:08.663621Z"
      }
    },
    "cell_type" : "code",
    "source" : "fn(x_test).round(2)",
    "id" : "965c488ed662fbf7",
    "outputs" : [ {
      "data" : {
        "text/plain" : [ "array([[ 1.56,  0.  ,  0.  ,  0.  ,  0.  ],\n", "       [ 4.6 ,  3.08,  0.  ,  0.  ,  0.  ],\n", "       [ 7.65,  9.  ,  5.8 ,  0.  ,  0.  ],\n", "       [11.54, 12.96, 15.99, 10.54,  0.  ],\n", "       [ 5.26,  6.07,  7.77, 10.03,  2.21]])" ]
      },
      "execution_count" : 17,
      "metadata" : { },
      "output_type" : "execute_result"
    } ],
    "execution_count" : 17
  }, {
    "metadata" : { },
    "cell_type" : "code",
    "outputs" : [ ],
    "execution_count" : null,
    "source" : "",
    "id" : "6edf35a6e4e87f58"
  } ],
  "metadata" : {
    "kernelspec" : {
      "display_name" : "Python 3",
      "language" : "python",
      "name" : "python3"
    },
    "language_info" : {
      "codemirror_mode" : {
        "name" : "ipython",
        "version" : 2
      },
      "file_extension" : ".py",
      "mimetype" : "text/x-python",
      "name" : "python",
      "nbconvert_exporter" : "python",
      "pygments_lexer" : "ipython2",
      "version" : "2.7.6"
    }
  },
  "nbformat" : 4,
  "nbformat_minor" : 5
}