lars/FYS-STK4155
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

change

Browse Source
This commit is contained in:
Morten Hjorth-Jensen
2023-10-15 21:46:13 +02:00
parent 295889032e
commit cb5ab447ad
2 changed files with 6 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/LectureNotes/_build/.doctrees/chapter11.doctree
BIN
View File
Binary file not shown.
doc/LectureNotes/_build/jupyter_execute/chapter11.ipynb
+6 -8
View File
@@ -3029,22 +3029,20 @@
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/wrap_util.py:20\u001b[0m, in \u001b[0;36munary_to_nary.<locals>.nary_operator.<locals>.nary_f\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 18\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 19\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(args[i] \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m argnum)\n\u001b[0;32m---> 20\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43munary_operator\u001b[49m\u001b[43m(\u001b[49m\u001b[43munary_f\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnary_op_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnary_op_kwargs\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/differential_operators.py:78\u001b[0m, in \u001b[0;36mhessian\u001b[0;34m(fun, x)\u001b[0m\n\u001b[1;32m 75\u001b[0m \u001b[38;5;129m@unary_to_nary\u001b[39m\n\u001b[1;32m 76\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mhessian\u001b[39m(fun, x):\n\u001b[1;32m 77\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mReturns a function that computes the exact Hessian.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m---> 78\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mjacobian\u001b[49m\u001b[43m(\u001b[49m\u001b[43mjacobian\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/wrap_util.py:20\u001b[0m, in \u001b[0;36munary_to_nary.<locals>.nary_operator.<locals>.nary_f\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 18\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 19\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(args[i] \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m argnum)\n\u001b[0;32m---> 20\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43munary_operator\u001b[49m\u001b[43m(\u001b[49m\u001b[43munary_f\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnary_op_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnary_op_kwargs\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/differential_operators.py:57\u001b[0m, in \u001b[0;36mjacobian\u001b[0;34m(fun, x)\u001b[0m\n\u001b[1;32m 47\u001b[0m \u001b[38;5;129m@unary_to_nary\u001b[39m\n\u001b[1;32m 48\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mjacobian\u001b[39m(fun, x):\n\u001b[1;32m 49\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 50\u001b[0m \u001b[38;5;124;03m Returns a function which computes the Jacobian of `fun` with respect to\u001b[39;00m\n\u001b[1;32m 51\u001b[0m \u001b[38;5;124;03m positional argument number `argnum`, which must be a scalar or array. Unlike\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 55\u001b[0m \u001b[38;5;124;03m (out1, out2, ...) then the Jacobian has shape (out1, out2, ..., in1, in2, ...).\u001b[39;00m\n\u001b[1;32m 56\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m---> 57\u001b[0m vjp, ans \u001b[38;5;241m=\u001b[39m \u001b[43m_make_vjp\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 58\u001b[0m ans_vspace \u001b[38;5;241m=\u001b[39m vspace(ans)\n\u001b[1;32m 59\u001b[0m jacobian_shape \u001b[38;5;241m=\u001b[39m ans_vspace\u001b[38;5;241m.\u001b[39mshape \u001b[38;5;241m+\u001b[39m vspace(x)\u001b[38;5;241m.\u001b[39mshape\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/core.py:10\u001b[0m, in \u001b[0;36mmake_vjp\u001b[0;34m(fun, x)\u001b[0m\n\u001b[1;32m 8\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmake_vjp\u001b[39m(fun, x):\n\u001b[1;32m 9\u001b[0m start_node \u001b[38;5;241m=\u001b[39m VJPNode\u001b[38;5;241m.\u001b[39mnew_root()\n\u001b[0;32m---> 10\u001b[0m end_value, end_node \u001b[38;5;241m=\u001b[39m \u001b[43mtrace\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstart_node\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 11\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m end_node \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 12\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvjp\u001b[39m(g): \u001b[38;5;28;01mreturn\u001b[39;00m vspace(x)\u001b[38;5;241m.\u001b[39mzeros()\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/tracer.py:10\u001b[0m, in \u001b[0;36mtrace\u001b[0;34m(start_node, fun, x)\u001b[0m\n\u001b[1;32m 8\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m trace_stack\u001b[38;5;241m.\u001b[39mnew_trace() \u001b[38;5;28;01mas\u001b[39;00m t:\n\u001b[1;32m 9\u001b[0m start_box \u001b[38;5;241m=\u001b[39m new_box(x, t, start_node)\n\u001b[0;32m---> 10\u001b[0m end_box \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstart_box\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 11\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m isbox(end_box) \u001b[38;5;129;01mand\u001b[39;00m end_box\u001b[38;5;241m.\u001b[39m_trace \u001b[38;5;241m==\u001b[39m start_box\u001b[38;5;241m.\u001b[39m_trace:\n\u001b[1;32m 12\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m end_box\u001b[38;5;241m.\u001b[39m_value, end_box\u001b[38;5;241m.\u001b[39m_node\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/wrap_util.py:15\u001b[0m, in \u001b[0;36munary_to_nary.<locals>.nary_operator.<locals>.nary_f.<locals>.unary_f\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 13\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 14\u001b[0m subargs \u001b[38;5;241m=\u001b[39m subvals(args, \u001b[38;5;28mzip\u001b[39m(argnum, x))\n\u001b[0;32m---> 15\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43msubargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/wrap_util.py:20\u001b[0m, in \u001b[0;36munary_to_nary.<locals>.nary_operator.<locals>.nary_f\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 18\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 19\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(args[i] \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m argnum)\n\u001b[0;32m---> 20\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43munary_operator\u001b[49m\u001b[43m(\u001b[49m\u001b[43munary_f\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnary_op_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnary_op_kwargs\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/differential_operators.py:61\u001b[0m, in \u001b[0;36mjacobian\u001b[0;34m(fun, x)\u001b[0m\n\u001b[1;32m 59\u001b[0m jacobian_shape \u001b[38;5;241m=\u001b[39m ans_vspace\u001b[38;5;241m.\u001b[39mshape \u001b[38;5;241m+\u001b[39m vspace(x)\u001b[38;5;241m.\u001b[39mshape\n\u001b[1;32m 60\u001b[0m grads \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(vjp, ans_vspace\u001b[38;5;241m.\u001b[39mstandard_basis())\n\u001b[0;32m---> 61\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m np\u001b[38;5;241m.\u001b[39mreshape(\u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstack\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgrads\u001b[49m\u001b[43m)\u001b[49m, jacobian_shape)\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/numpy/numpy_wrapper.py:88\u001b[0m, in \u001b[0;36mstack\u001b[0;34m(arrays, axis)\u001b[0m\n\u001b[1;32m 83\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstack\u001b[39m(arrays, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m):\n\u001b[1;32m 84\u001b[0m \u001b[38;5;66;03m# this code is basically copied from numpy/core/shape_base.py's stack\u001b[39;00m\n\u001b[1;32m 85\u001b[0m \u001b[38;5;66;03m# we need it here because we want to re-implement stack in terms of the\u001b[39;00m\n\u001b[1;32m 86\u001b[0m \u001b[38;5;66;03m# primitives defined in this file\u001b[39;00m\n\u001b[0;32m---> 88\u001b[0m arrays \u001b[38;5;241m=\u001b[39m [array(arr) \u001b[38;5;28;01mfor\u001b[39;00m arr \u001b[38;5;129;01min\u001b[39;00m arrays]\n\u001b[1;32m 89\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m arrays:\n\u001b[1;32m 90\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mneed at least one array to stack\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/numpy/numpy_wrapper.py:88\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 83\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstack\u001b[39m(arrays, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m):\n\u001b[1;32m 84\u001b[0m \u001b[38;5;66;03m# this code is basically copied from numpy/core/shape_base.py's stack\u001b[39;00m\n\u001b[1;32m 85\u001b[0m \u001b[38;5;66;03m# we need it here because we want to re-implement stack in terms of the\u001b[39;00m\n\u001b[1;32m 86\u001b[0m \u001b[38;5;66;03m# primitives defined in this file\u001b[39;00m\n\u001b[0;32m---> 88\u001b[0m arrays \u001b[38;5;241m=\u001b[39m [array(arr) \u001b[38;5;28;01mfor\u001b[39;00m arr \u001b[38;5;129;01min\u001b[39;00m arrays]\n\u001b[1;32m 89\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m arrays:\n\u001b[1;32m 90\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mneed at least one array to stack\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/core.py:14\u001b[0m, in \u001b[0;36mmake_vjp.<locals>.vjp\u001b[0;34m(g)\u001b[0m\n\u001b[0;32m---> 14\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvjp\u001b[39m(g): \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mbackward_pass\u001b[49m\u001b[43m(\u001b[49m\u001b[43mg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mend_node\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/core.py:21\u001b[0m, in \u001b[0;36mbackward_pass\u001b[0;34m(g, end_node)\u001b[0m\n\u001b[1;32m 19\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m node \u001b[38;5;129;01min\u001b[39;00m toposort(end_node):\n\u001b[1;32m 20\u001b[0m outgrad \u001b[38;5;241m=\u001b[39m outgrads\u001b[38;5;241m.\u001b[39mpop(node)\n\u001b[0;32m---> 21\u001b[0m ingrads \u001b[38;5;241m=\u001b[39m \u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvjp\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutgrad\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 22\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m parent, ingrad \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(node\u001b[38;5;241m.\u001b[39mparents, ingrads):\n\u001b[1;32m 23\u001b[0m outgrads[parent] \u001b[38;5;241m=\u001b[39m add_outgrads(outgrads\u001b[38;5;241m.\u001b[39mget(parent), ingrad)\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/core.py:67\u001b[0m, in \u001b[0;36mdefvjp.<locals>.vjp_argnums.<locals>.<lambda>\u001b[0;34m(g)\u001b[0m\n\u001b[1;32m 64\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(\n\u001b[1;32m 65\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mVJP of \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m wrt argnum 0 not defined\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(fun\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m))\n\u001b[1;32m 66\u001b[0m vjp \u001b[38;5;241m=\u001b[39m vjpfun(ans, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 67\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mlambda\u001b[39;00m g: (\u001b[43mvjp\u001b[49m\u001b[43m(\u001b[49m\u001b[43mg\u001b[49m\u001b[43m)\u001b[49m,)\n\u001b[1;32m 68\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m L \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m2\u001b[39m:\n\u001b[1;32m 69\u001b[0m argnum_0, argnum_1 \u001b[38;5;241m=\u001b[39m argnums\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/numpy/numpy_vjps.py:701\u001b[0m, in \u001b[0;36m<lambda>\u001b[0;34m(g)\u001b[0m\n\u001b[1;32m 699\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m A\n\u001b[1;32m 700\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m SparseObject(vs, mut_add)\n\u001b[0;32m--> 701\u001b[0m defvjp(func(ArrayBox\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__getitem__\u001b[39m), \u001b[38;5;28;01mlambda\u001b[39;00m ans, A, idx: \u001b[38;5;28;01mlambda\u001b[39;00m g: \u001b[43muntake\u001b[49m\u001b[43m(\u001b[49m\u001b[43mg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43midx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvspace\u001b[49m\u001b[43m(\u001b[49m\u001b[43mA\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m 702\u001b[0m defvjp(untake, \u001b[38;5;28;01mlambda\u001b[39;00m ans, x, idx, _: \u001b[38;5;28;01mlambda\u001b[39;00m g: g[idx])\n\u001b[1;32m 704\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_unpad\u001b[39m(array, width):\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/tracer.py:44\u001b[0m, in \u001b[0;36mprimitive.<locals>.f_wrapped\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 42\u001b[0m parents \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(box\u001b[38;5;241m.\u001b[39m_node \u001b[38;5;28;01mfor\u001b[39;00m _ , box \u001b[38;5;129;01min\u001b[39;00m boxed_args)\n\u001b[1;32m 43\u001b[0m argnums \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(argnum \u001b[38;5;28;01mfor\u001b[39;00m argnum, _ \u001b[38;5;129;01min\u001b[39;00m boxed_args)\n\u001b[0;32m---> 44\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mf_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margvals\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 45\u001b[0m node \u001b[38;5;241m=\u001b[39m node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)\n\u001b[1;32m 46\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m new_box(ans, trace, node)\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/numpy/numpy_vjps.py:423\u001b[0m, in \u001b[0;36mmatmul_vjp_1.<locals>.<lambda>\u001b[0;34m(g)\u001b[0m\n\u001b[1;32m 421\u001b[0m A_ndim \u001b[38;5;241m=\u001b[39m anp\u001b[38;5;241m.\u001b[39mndim(A)\n\u001b[1;32m 422\u001b[0m B_meta \u001b[38;5;241m=\u001b[39m anp\u001b[38;5;241m.\u001b[39mmetadata(B)\n\u001b[0;32m--> 423\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mlambda\u001b[39;00m g: \u001b[43mmatmul_adjoint_1\u001b[49m\u001b[43m(\u001b[49m\u001b[43mA\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mA_ndim\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mB_meta\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/numpy/numpy_vjps.py:410\u001b[0m, in \u001b[0;36mmatmul_adjoint_1\u001b[0;34m(A, G, A_ndim, B_meta)\u001b[0m\n\u001b[1;32m 408\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m: \u001b[38;5;66;03m# We need to swap the last two axes of A\u001b[39;00m\n\u001b[1;32m 409\u001b[0m A \u001b[38;5;241m=\u001b[39m anp\u001b[38;5;241m.\u001b[39mswapaxes(A, A_ndim \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m2\u001b[39m, A_ndim \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m--> 410\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43manp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmatmul\u001b[49m\u001b[43m(\u001b[49m\u001b[43mA\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mG\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 411\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m B_is_vec:\n\u001b[1;32m 412\u001b[0m result \u001b[38;5;241m=\u001b[39m anp\u001b[38;5;241m.\u001b[39msqueeze(result, anp\u001b[38;5;241m.\u001b[39mndim(G) \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m)\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/tracer.py:45\u001b[0m, in \u001b[0;36mprimitive.<locals>.f_wrapped\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 43\u001b[0m argnums \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(argnum \u001b[38;5;28;01mfor\u001b[39;00m argnum, _ \u001b[38;5;129;01min\u001b[39;00m boxed_args)\n\u001b[1;32m 44\u001b[0m ans \u001b[38;5;241m=\u001b[39m f_wrapped(\u001b[38;5;241m*\u001b[39margvals, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 45\u001b[0m node \u001b[38;5;241m=\u001b[39m \u001b[43mnode_constructor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mans\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mf_wrapped\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margvals\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margnums\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparents\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 46\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m new_box(ans, trace, node)\n\u001b[1;32m 47\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/core.py:36\u001b[0m, in \u001b[0;36mVJPNode.__init__\u001b[0;34m(self, value, fun, args, kwargs, parent_argnums, parents)\u001b[0m\n\u001b[1;32m 33\u001b[0m fun_name \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(fun, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__name__\u001b[39m\u001b[38;5;124m'\u001b[39m, fun)\n\u001b[1;32m 34\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mVJP of \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m wrt argnums \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m not defined\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;241m.\u001b[39mformat(fun_name, parent_argnums))\n\u001b[0;32m---> 36\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvjp \u001b[38;5;241m=\u001b[39m \u001b[43mvjpmaker\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparent_argnums\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/core.py:56\u001b[0m, in \u001b[0;36mdefvjp.<locals>.vjp_argnums\u001b[0;34m(argnums, ans, args, kwargs)\u001b[0m\n\u001b[1;32m 53\u001b[0m argnums \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124margnums\u001b[39m\u001b[38;5;124m'\u001b[39m, count())\n\u001b[1;32m 54\u001b[0m vjps_dict \u001b[38;5;241m=\u001b[39m {argnum : translate_vjp(vjpmaker, fun, argnum)\n\u001b[1;32m 55\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m argnum, vjpmaker \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(argnums, vjpmakers)}\n\u001b[0;32m---> 56\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvjp_argnums\u001b[39m(argnums, ans, args, kwargs):\n\u001b[1;32m 57\u001b[0m L \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(argnums)\n\u001b[1;32m 58\u001b[0m \u001b[38;5;66;03m# These first two cases are just optimizations\u001b[39;00m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/core.py:77\u001b[0m, in \u001b[0;36mdefvjp.<locals>.vjp_argnums\u001b[0;34m(argnums, ans, args, kwargs)\u001b[0m\n\u001b[1;32m 74\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(\n\u001b[1;32m 75\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mVJP of \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m wrt argnums 0, 1 not defined\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(fun\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m))\n\u001b[1;32m 76\u001b[0m vjp_0 \u001b[38;5;241m=\u001b[39m vjp_0_fun(ans, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 77\u001b[0m vjp_1 \u001b[38;5;241m=\u001b[39m \u001b[43mvjp_1_fun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mans\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 78\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mlambda\u001b[39;00m g: (vjp_0(g), vjp_1(g))\n\u001b[1;32m 79\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/numpy/numpy_vjps.py:422\u001b[0m, in \u001b[0;36mmatmul_vjp_1\u001b[0;34m(ans, A, B)\u001b[0m\n\u001b[1;32m 420\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmatmul_vjp_1\u001b[39m(ans, A, B):\n\u001b[1;32m 421\u001b[0m A_ndim \u001b[38;5;241m=\u001b[39m anp\u001b[38;5;241m.\u001b[39mndim(A)\n\u001b[0;32m--> 422\u001b[0m B_meta \u001b[38;5;241m=\u001b[39m \u001b[43manp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmetadata\u001b[49m\u001b[43m(\u001b[49m\u001b[43mB\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 423\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mlambda\u001b[39;00m g: matmul_adjoint_1(A, g, A_ndim, B_meta)\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/tracer.py:61\u001b[0m, in \u001b[0;36mnotrace_primitive.<locals>.f_wrapped\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 58\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(f_raw)\n\u001b[1;32m 59\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mf_wrapped\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m 60\u001b[0m argvals \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(getval, args)\n\u001b[0;32m---> 61\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mf_raw\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margvals\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/autograd/numpy/numpy_wrapper.py:148\u001b[0m, in \u001b[0;36mmetadata\u001b[0;34m(A)\u001b[0m\n\u001b[1;32m 146\u001b[0m \u001b[38;5;129m@notrace_primitive\u001b[39m\n\u001b[1;32m 147\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmetadata\u001b[39m(A):\n\u001b[0;32m--> 148\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _np\u001b[38;5;241m.\u001b[39mshape(A), _np\u001b[38;5;241m.\u001b[39mndim(A), _np\u001b[38;5;241m.\u001b[39mresult_type(A), _np\u001b[38;5;241m.\u001b[39miscomplexobj(A)\n",
"\u001b[0;31mKeyboardInterrupt\u001b[0m: "
]
}