From cf175f5393c8b513a4e9a698cfe56bb5e8710c0a Mon Sep 17 00:00:00 2001 From: Grace Perkins Date: Thu, 15 Oct 2026 17:41:46 -0400 Subject: [PATCH] Added chemical potentials mu_nue and mu_numt and their dmu equations --- linx/Untitled.ipynb | 152 +++++++++++++++++- linx/background.py | 37 +++-- linx/chemical potential testing.ipynb | 2 +- ...linx chemical potential testing 4-30.ipynb | 4 +- 4 files changed, 175 insertions(+), 20 deletions(-) diff --git a/linx/Untitled.ipynb b/linx/Untitled.ipynb index 125abfb..ab34a7a 100644 --- a/linx/Untitled.ipynb +++ b/linx/Untitled.ipynb @@ -2,20 +2,164 @@ "cells": [ { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "a0074e73-6ca5-47ee-a7bc-87a25e8fb3a5", "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "`\\ /´ |||| |||| ||||| |||| |||| ||||\n", + " /\\_______/\\ |||| |||| ||||||| |||| |||| ||||\n", + " ) __` ´__ ( |||| |||| |||| |||| |||| |||||||\n", + "/ `-|_|-´ \\ |||| |||| |||| |||| ||| ||||||| \n", + "/ (_x_) \\ |||||||||| |||| |||| ||||||| |||| ||||\n", + " ) `-´ ( |||||||||| |||| |||| |||||| |||| ||||\n", + " \n", + "Compiling thermodynamics model...\n" + ] + }, + { + "ename": "KeyboardInterrupt", + "evalue": "", + "output_type": "error", + "traceback": [ + "\u001b[31m---------------------------------------------------------------------------\u001b[39m", + "\u001b[31mJaxStackTraceBeforeTransformation\u001b[39m Traceback (most recent call last)", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\internal\\_loop\\checkpointed.py:271\u001b[39m, in \u001b[36m_checkpointed_while_loop\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m 269\u001b[39m \u001b[38;5;66;03m# Hashable wrapper; JAX issue #13554 and\u001b[39;00m\n\u001b[32m 270\u001b[39m \u001b[38;5;66;03m# https://github.com/patrick-kidger/equinox/issues/768\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m271\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m while_loop(\u001b[38;5;28;01mlambda\u001b[39;00m x: cond_fun(x), \u001b[38;5;28;01mlambda\u001b[39;00m x: body_fun(x), init_val)\n", + "\u001b[31mJaxStackTraceBeforeTransformation\u001b[39m: KeyboardInterrupt\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------", + "\nThe above exception was the direct cause of the following exception:\n", + "\u001b[31mKeyboardInterrupt\u001b[39m Traceback (most recent call last)", + "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 21\u001b[39m\n\u001b[32m 17\u001b[39m thermo_model_DNeff = BackgroundModel(max_steps=\u001b[32m1\u001b[39m)\n\u001b[32m 18\u001b[39m \n\u001b[32m 19\u001b[39m (\n\u001b[32m 20\u001b[39m t_vec_ref, a_vec_ref, rho_g_vec, rho_nu_vec, rho_NP_vec, P_NP_vec, Neff_vec\n\u001b[32m---> \u001b[39m\u001b[32m21\u001b[39m ) = thermo_model_DNeff(0.)\n\u001b[32m 22\u001b[39m \n\u001b[32m 23\u001b[39m network = \u001b[33m'key_PRIMAT_2023'\u001b[39m\n\u001b[32m 24\u001b[39m \u001b[38;5;66;03m# network = 'key_PRIMAT_2018'\u001b[39;00m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\_module\\_prebuilt.py:82\u001b[39m, in \u001b[36mPartial.__call__\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 69\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args: Any, **kwargs: Any) -> _Return:\n\u001b[32m 70\u001b[39m \u001b[38;5;250m \u001b[39m\u001b[33;03m\"\"\"Call the wrapped `self.func`.\u001b[39;00m\n\u001b[32m 71\u001b[39m \n\u001b[32m 72\u001b[39m \u001b[33;03m **Arguments:**\u001b[39;00m\n\u001b[32m (...)\u001b[39m\u001b[32m 80\u001b[39m \u001b[33;03m The result of the wrapped function.\u001b[39;00m\n\u001b[32m 81\u001b[39m \u001b[33;03m \"\"\"\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m82\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mfunc\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mkeywords\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + " \u001b[31m[... skipping hidden 2 frame]\u001b[39m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\_jit.py:55\u001b[39m, in \u001b[36m_filter_jit_cache..fun_wrapped\u001b[39m\u001b[34m(dynamic_donate, dynamic_nodonate, static)\u001b[39m\n\u001b[32m 53\u001b[39m *args, dummy_arg = (first_arg,) + rest_args\n\u001b[32m 54\u001b[39m \u001b[38;5;28;01massert\u001b[39;00m dummy_arg \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m55\u001b[39m out = \u001b[30;43mfun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 56\u001b[39m dynamic_out, static_out = partition(out, is_array)\n\u001b[32m 57\u001b[39m marker = jnp.array(\u001b[32m0\u001b[39m)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\LINX-original\\linx\\..\\linx\\background.py:154\u001b[39m, in \u001b[36mBackgroundModel.__call__\u001b[39m\u001b[34m(self, Delt_Neff_init, T_start, T_end, me, rtol, atol, solver)\u001b[39m\n\u001b[32m 151\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mT_EM_check\u001b[39m(t, y, args, **kwargs): \n\u001b[32m 152\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m y[\u001b[32m1\u001b[39m] < T_end\n\u001b[32m--> \u001b[39m\u001b[32m154\u001b[39m sol = \u001b[30;43mdiffeqsolve\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 155\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mODETerm\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mdY\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43msolver\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mlna_init\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mrho_extra_init\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mme\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 156\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mt0\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mt0\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mt1\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mjnp\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43minf\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mdt0\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43;01mNone\u001b[39;49;00m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43my0\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mY0\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\n\u001b[32m 157\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43msaveat\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mSaveAt\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43msteps\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43;01mTrue\u001b[39;49;00m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mevent\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mEvent\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mT_EM_check\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 158\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mstepsize_controller\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mPIDController\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 159\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mrtol\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mrtol\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43matol\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43matol\u001b[39;49m\n\u001b[32m 160\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\n\u001b[32m 161\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 162\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mthrow\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mthrow\u001b[39;49m\n\u001b[32m 163\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 165\u001b[39m a_vec = jnp.exp(sol.ys[\u001b[32m0\u001b[39m])\n\u001b[32m 166\u001b[39m rho_g_vec = rho_massless_BE_v(sol.ys[\u001b[32m1\u001b[39m], \u001b[32m0.\u001b[39m, \u001b[32m2\u001b[39m) \n", + " \u001b[31m[... skipping hidden 2 frame]\u001b[39m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\_jit.py:55\u001b[39m, in \u001b[36m_filter_jit_cache..fun_wrapped\u001b[39m\u001b[34m(dynamic_donate, dynamic_nodonate, static)\u001b[39m\n\u001b[32m 53\u001b[39m *args, dummy_arg = (first_arg,) + rest_args\n\u001b[32m 54\u001b[39m \u001b[38;5;28;01massert\u001b[39;00m dummy_arg \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m55\u001b[39m out = \u001b[30;43mfun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 56\u001b[39m dynamic_out, static_out = partition(out, is_array)\n\u001b[32m 57\u001b[39m marker = jnp.array(\u001b[32m0\u001b[39m)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\diffrax\\_integrate.py:1456\u001b[39m, in \u001b[36mdiffeqsolve\u001b[39m\u001b[34m(terms, solver, t0, t1, dt0, y0, args, saveat, stepsize_controller, adjoint, event, max_steps, throw, progress_meter, solver_state, controller_state, made_jump, discrete_terminating_event)\u001b[39m\n\u001b[32m 1429\u001b[39m init_state = State(\n\u001b[32m 1430\u001b[39m y=y0,\n\u001b[32m 1431\u001b[39m tprev=tprev,\n\u001b[32m (...)\u001b[39m\u001b[32m 1449\u001b[39m event_mask=event_mask,\n\u001b[32m 1450\u001b[39m )\n\u001b[32m 1452\u001b[39m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[32m 1453\u001b[39m \u001b[38;5;66;03m# Main loop\u001b[39;00m\n\u001b[32m 1454\u001b[39m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m1456\u001b[39m final_state, aux_stats = \u001b[30;43madjoint\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mloop\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 1457\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1458\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mterms\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mterms\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1459\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43msolver\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43msolver\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1460\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mstepsize_controller\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mstepsize_controller\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1461\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mevent\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mevent\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1462\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43msaveat\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43msaveat\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1463\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mt0\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mt0\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1464\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mt1\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mt1\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1465\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mdt0\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mdt0\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1466\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1467\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43minit_state\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43minit_state\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1468\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mthrow\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mthrow\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1469\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mpassed_solver_state\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mpassed_solver_state\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1470\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mpassed_controller_state\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mpassed_controller_state\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1471\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mprogress_meter\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mprogress_meter\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1472\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1474\u001b[39m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[32m 1475\u001b[39m \u001b[38;5;66;03m# Finish up\u001b[39;00m\n\u001b[32m 1476\u001b[39m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[32m 1478\u001b[39m progress_meter.close(final_state.progress_meter_state)\n", + " \u001b[31m[... skipping hidden 1 frame]\u001b[39m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\diffrax\\_adjoint.py:299\u001b[39m, in \u001b[36mRecursiveCheckpointAdjoint.loop\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m 295\u001b[39m outer_while_loop = ft.partial(\n\u001b[32m 296\u001b[39m _outer_loop, kind=\u001b[33m\"\u001b[39m\u001b[33mcheckpointed\u001b[39m\u001b[33m\"\u001b[39m, checkpoints=\u001b[38;5;28mself\u001b[39m.checkpoints\n\u001b[32m 297\u001b[39m )\n\u001b[32m 298\u001b[39m msg = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m299\u001b[39m final_state = \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43m_loop\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 300\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mterms\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mterms\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 301\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43msaveat\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43msaveat\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 302\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43minit_state\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43minit_state\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 303\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 304\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43minner_while_loop\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43minner_while_loop\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 305\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mouter_while_loop\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mouter_while_loop\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 306\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 307\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 308\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m msg \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m 309\u001b[39m final_state = eqxi.nondifferentiable_backward(\n\u001b[32m 310\u001b[39m final_state, msg=msg, symbolic=\u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[32m 311\u001b[39m )\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\diffrax\\_integrate.py:685\u001b[39m, in \u001b[36mloop\u001b[39m\u001b[34m(solver, stepsize_controller, event, saveat, t0, t1, dt0, max_steps, terms, args, init_state, inner_while_loop, outer_while_loop, progress_meter)\u001b[39m\n\u001b[32m 682\u001b[39m new_state, _, _ = body_fun_aux(state)\n\u001b[32m 683\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m new_state\n\u001b[32m--> \u001b[39m\u001b[32m685\u001b[39m final_state = \u001b[30;43mouter_while_loop\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 686\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcond_fun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mbody_fun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43minit_state\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mbuffers\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43m_outer_buffers\u001b[39;49m\n\u001b[32m 687\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 688\u001b[39m save_state = final_state.save_state\n\u001b[32m 689\u001b[39m result = final_state.result\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\source_info_util.py:291\u001b[39m, in \u001b[36mExtendNameStackContextManager.__call__..wrapper\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m 288\u001b[39m \u001b[38;5;129m@functools\u001b[39m.wraps(func)\n\u001b[32m 289\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mwrapper\u001b[39m(*args, **kwargs):\n\u001b[32m 290\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m ExtendNameStackContextManager(\u001b[38;5;28mself\u001b[39m.name):\n\u001b[32m--> \u001b[39m\u001b[32m291\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mfunc\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\internal\\_loop\\loop.py:107\u001b[39m, in \u001b[36mwhile_loop\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m 105\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m kind == \u001b[33m\"\u001b[39m\u001b[33mcheckpointed\u001b[39m\u001b[33m\"\u001b[39m:\n\u001b[32m 106\u001b[39m \u001b[38;5;28;01mdel\u001b[39;00m kind, base\n\u001b[32m--> \u001b[39m\u001b[32m107\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mcheckpointed_while_loop\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 108\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcond_fun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 109\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mbody_fun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 110\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43minit_val\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 111\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 112\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mbuffers\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mbuffers\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 113\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcheckpoints\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mcheckpoints\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 114\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 115\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m kind == \u001b[33m\"\u001b[39m\u001b[33mbounded\u001b[39m\u001b[33m\"\u001b[39m:\n\u001b[32m 116\u001b[39m \u001b[38;5;28;01mdel\u001b[39;00m kind, checkpoints\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\internal\\_loop\\checkpointed.py:249\u001b[39m, in \u001b[36mcheckpointed_while_loop\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m 247\u001b[39m body_fun_ = filter_closure_convert(body_fun_, init_val_)\n\u001b[32m 248\u001b[39m vjp_arg = (init_val_, body_fun_)\n\u001b[32m--> \u001b[39m\u001b[32m249\u001b[39m final_val_ = \u001b[30;43m_checkpointed_while_loop\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 250\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mvjp_arg\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcond_fun_\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcheckpoints\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mbuffers_\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mmax_steps\u001b[39;49m\n\u001b[32m 251\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 252\u001b[39m _, _, _, final_val = _stop_gradient_on_unperturbed(init_val_, final_val_, body_fun_)\n\u001b[32m 253\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m final_val\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\_ad.py:1107\u001b[39m, in \u001b[36mfilter_custom_vjp.__call__\u001b[39m\u001b[34m(self, vjp_arg, *args, **kwargs)\u001b[39m\n\u001b[32m 1103\u001b[39m array_args_kwargs, nonarray_args_kwargs = partition((args, kwargs), is_array)\n\u001b[32m 1104\u001b[39m array_args_kwargs = nondifferentiable(\n\u001b[32m 1105\u001b[39m array_args_kwargs, name=\u001b[33m\"\u001b[39m\u001b[33m`*args` and `**kwargs` to `filter_custom_vjp`\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 1106\u001b[39m )\n\u001b[32m-> \u001b[39m\u001b[32m1107\u001b[39m out = \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mfn_wrapped\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 1108\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mnonarray_vjp_arg\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1109\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mnonarray_args_kwargs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1110\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mdiff_array_vjp_arg\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1111\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mnondiff_array_vjp_arg\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1112\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43marray_args_kwargs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1113\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1114\u001b[39m diff_array_out, nondiff_array_out, nonarray_out = out\n\u001b[32m 1115\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m combine(diff_array_out, nondiff_array_out, nonarray_out.value)\n", + " \u001b[31m[... skipping hidden 1 frame]\u001b[39m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\custom_derivatives.py:771\u001b[39m, in \u001b[36mcustom_vjp.__call__\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 767\u001b[39m flat_fwd, out_trees = _flatten_fwd(\n\u001b[32m 768\u001b[39m fwd_, \u001b[38;5;28mself\u001b[39m.nondiff_argnums, \u001b[38;5;28mself\u001b[39m.symbolic_zeros, debug_fun,\n\u001b[32m 769\u001b[39m debug_fwd, in_tree, out_type)\n\u001b[32m 770\u001b[39m flat_bwd = _flatten_bwd(bwd, in_tree, in_avals, out_trees)\n\u001b[32m--> \u001b[39m\u001b[32m771\u001b[39m out_flat = \u001b[30;43mcustom_vjp_call_p\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mbind\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs_flat\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43msubfuns\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mflat_fun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mflat_fwd\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mflat_bwd\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 772\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mout_trees\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mout_trees\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 773\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43msymbolic_zeros\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43msymbolic_zeros\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 774\u001b[39m _, (out_tree, _, _) = lu.merge_linear_aux(out_type, out_trees)\n\u001b[32m 775\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m tree_unflatten(out_tree, out_flat)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:705\u001b[39m, in \u001b[36mPrimitive.bind\u001b[39m\u001b[34m(self, *args, **params)\u001b[39m\n\u001b[32m 703\u001b[39m trace_ctx.set_trace(eval_trace)\n\u001b[32m 704\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m705\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mbind_with_trace\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mprev_trace\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mavals\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 706\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m 707\u001b[39m trace_ctx.set_trace(prev_trace)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\custom_derivatives.py:1023\u001b[39m, in \u001b[36mCustomVJPCallPrimitive.bind_with_trace\u001b[39m\u001b[34m(self, trace, args, avals, params)\u001b[39m\n\u001b[32m 1021\u001b[39m params = \u001b[38;5;28mdict\u001b[39m(params)\n\u001b[32m 1022\u001b[39m fun, fwd, bwd = params.pop(\u001b[33m'\u001b[39m\u001b[33msubfuns\u001b[39m\u001b[33m'\u001b[39m)\n\u001b[32m-> \u001b[39m\u001b[32m1023\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mtrace\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mprocess_custom_vjp_call\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mfun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mfwd\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mbwd\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:1310\u001b[39m, in \u001b[36mEvalTrace.process_custom_vjp_call\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m 1308\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mprocess_custom_vjp_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, fun, fwd, bwd, tracers, /, **_):\n\u001b[32m 1309\u001b[39m \u001b[38;5;28;01mdel\u001b[39;00m primitive, fwd, bwd, _ \u001b[38;5;66;03m# Unused.\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m1310\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mfun\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcall_wrapped\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mtracers\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\linear_util.py:212\u001b[39m, in \u001b[36mWrappedFun.call_wrapped\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 210\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mcall_wrapped\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args, **kwargs):\n\u001b[32m 211\u001b[39m \u001b[38;5;250m \u001b[39m\u001b[33;03m\"\"\"Calls the transformed function\"\"\"\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m212\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mf_transformed\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\custom_derivatives.py:108\u001b[39m, in \u001b[36m_flatten_fun_nokwargs\u001b[39m\u001b[34m(f, store, in_tree, *args_flat)\u001b[39m\n\u001b[32m 103\u001b[39m \u001b[38;5;129m@lu\u001b[39m.transformation_with_aux2\n\u001b[32m 104\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_flatten_fun_nokwargs\u001b[39m(f: Callable,\n\u001b[32m 105\u001b[39m store: lu.Store, in_tree: PyTreeDef,\n\u001b[32m 106\u001b[39m *args_flat):\n\u001b[32m 107\u001b[39m py_args = tree_unflatten(in_tree, args_flat)\n\u001b[32m--> \u001b[39m\u001b[32m108\u001b[39m ans = \u001b[30;43mf\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mpy_args\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 109\u001b[39m ans_flat, ans_tree = tree_flatten(ans)\n\u001b[32m 110\u001b[39m ans_avals = [core.typeof(x) \u001b[38;5;28;01mfor\u001b[39;00m x \u001b[38;5;129;01min\u001b[39;00m ans_flat]\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\custom_derivatives.py:781\u001b[39m, in \u001b[36m_check_primal_refs\u001b[39m\u001b[34m(f, nondiff_argnums, debug, *args)\u001b[39m\n\u001b[32m 777\u001b[39m \u001b[38;5;129m@lu\u001b[39m.transformation2\n\u001b[32m 778\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_check_primal_refs\u001b[39m(\n\u001b[32m 779\u001b[39m f: Callable, nondiff_argnums: Sequence[\u001b[38;5;28mint\u001b[39m], debug: core.DebugInfo, *args):\n\u001b[32m 780\u001b[39m _check_for_aliased_refs(f, nondiff_argnums, debug, args)\n\u001b[32m--> \u001b[39m\u001b[32m781\u001b[39m out = \u001b[30;43mf\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 782\u001b[39m _check_for_returned_refs(f, out, \u001b[33m'\u001b[39m\u001b[33mprimal\u001b[39m\u001b[33m'\u001b[39m, [], \u001b[32m0\u001b[39m)\n\u001b[32m 783\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m out\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\api_util.py:254\u001b[39m, in \u001b[36m_argnums_partial\u001b[39m\u001b[34m(_fun, _dyn_argnums, _fixed_args, *dyn_args, **kwargs)\u001b[39m\n\u001b[32m 252\u001b[39m args = [\u001b[38;5;28mnext\u001b[39m(fixed_args_).val \u001b[38;5;28;01mif\u001b[39;00m x \u001b[38;5;129;01mis\u001b[39;00m sentinel \u001b[38;5;28;01melse\u001b[39;00m x \u001b[38;5;28;01mfor\u001b[39;00m x \u001b[38;5;129;01min\u001b[39;00m args]\n\u001b[32m 253\u001b[39m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mnext\u001b[39m(fixed_args_, sentinel) \u001b[38;5;129;01mis\u001b[39;00m sentinel\n\u001b[32m--> \u001b[39m\u001b[32m254\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43m_fun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\_ad.py:1012\u001b[39m, in \u001b[36mfilter_custom_vjp._defvjp..fn_wrapped\u001b[39m\u001b[34m(nonarray_vjp_arg, nonarray_args_kwargs, diff_array_vjp_arg, nondiff_array_vjp_arg, array_args_kwargs)\u001b[39m\n\u001b[32m 1008\u001b[39m vjp_arg = combine(\n\u001b[32m 1009\u001b[39m nonarray_vjp_arg, diff_array_vjp_arg, nondiff_array_vjp_arg\n\u001b[32m 1010\u001b[39m )\n\u001b[32m 1011\u001b[39m args, kwargs = combine(nonarray_args_kwargs, array_args_kwargs)\n\u001b[32m-> \u001b[39m\u001b[32m1012\u001b[39m out = \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mfn\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mvjp_arg\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1013\u001b[39m array_out, nonarray_out = partition(out, is_array)\n\u001b[32m 1014\u001b[39m diff_array_out, nondiff_array_out = partition(array_out, is_inexact_array)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\internal\\_loop\\checkpointed.py:271\u001b[39m, in \u001b[36m_checkpointed_while_loop\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m 268\u001b[39m while_loop = jax.named_call(lax.while_loop, name=\u001b[33m\"\u001b[39m\u001b[33mcheckpointed-no-vjp\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 269\u001b[39m \u001b[38;5;66;03m# Hashable wrapper; JAX issue #13554 and\u001b[39;00m\n\u001b[32m 270\u001b[39m \u001b[38;5;66;03m# https://github.com/patrick-kidger/equinox/issues/768\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m271\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mwhile_loop\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43;01mlambda\u001b[39;49;00m\u001b[30;43m \u001b[39;49m\u001b[30;43mx\u001b[39;49m\u001b[30;43m:\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcond_fun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mx\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43;01mlambda\u001b[39;49;00m\u001b[30;43m \u001b[39;49m\u001b[30;43mx\u001b[39;49m\u001b[30;43m:\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mbody_fun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mx\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43minit_val\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\source_info_util.py:291\u001b[39m, in \u001b[36mExtendNameStackContextManager.__call__..wrapper\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m 288\u001b[39m \u001b[38;5;129m@functools\u001b[39m.wraps(func)\n\u001b[32m 289\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mwrapper\u001b[39m(*args, **kwargs):\n\u001b[32m 290\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m ExtendNameStackContextManager(\u001b[38;5;28mself\u001b[39m.name):\n\u001b[32m--> \u001b[39m\u001b[32m291\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mfunc\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + " \u001b[31m[... skipping hidden 1 frame]\u001b[39m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\lax\\control_flow\\loops.py:1605\u001b[39m, in \u001b[36mwhile_loop\u001b[39m\u001b[34m(cond_fun, body_fun, init_val)\u001b[39m\n\u001b[32m 1603\u001b[39m val = tree_map(lax.asarray, init_val)\n\u001b[32m 1604\u001b[39m \u001b[38;5;28;01mwhile\u001b[39;00m cond_fun(val):\n\u001b[32m-> \u001b[39m\u001b[32m1605\u001b[39m val = tree_map(lax.asarray, \u001b[30;43mbody_fun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mval\u001b[39;49m\u001b[30;43m)\u001b[39;49m)\n\u001b[32m 1606\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m val\n\u001b[32m 1607\u001b[39m \u001b[38;5;28;01mexcept\u001b[39;00m core.ConcretizationTypeError:\n\u001b[32m 1608\u001b[39m \u001b[38;5;66;03m# Can't run this while_loop in Python (e.g. because there's a vmap\u001b[39;00m\n\u001b[32m 1609\u001b[39m \u001b[38;5;66;03m# transformation on it), so we fall back to the primitive version.\u001b[39;00m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\internal\\_loop\\checkpointed.py:271\u001b[39m, in \u001b[36m_checkpointed_while_loop..\u001b[39m\u001b[34m(x)\u001b[39m\n\u001b[32m 268\u001b[39m while_loop = jax.named_call(lax.while_loop, name=\u001b[33m\"\u001b[39m\u001b[33mcheckpointed-no-vjp\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 269\u001b[39m \u001b[38;5;66;03m# Hashable wrapper; JAX issue #13554 and\u001b[39;00m\n\u001b[32m 270\u001b[39m \u001b[38;5;66;03m# https://github.com/patrick-kidger/equinox/issues/768\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m271\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m while_loop(\u001b[38;5;28;01mlambda\u001b[39;00m x: cond_fun(x), \u001b[38;5;28;01mlambda\u001b[39;00m x: \u001b[30;43mbody_fun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mx\u001b[39;49m\u001b[30;43m)\u001b[39;49m, init_val)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\equinox\\_ad.py:618\u001b[39m, in \u001b[36m_ClosureConvert.__call__\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 616\u001b[39m in_dynamic = _check_closure_convert_input(\u001b[38;5;28mself\u001b[39m, args, kwargs)\n\u001b[32m 617\u001b[39m in_dynamic_flat = jtu.tree_leaves(in_dynamic)\n\u001b[32m--> \u001b[39m\u001b[32m618\u001b[39m out_dynamic_flat = \u001b[30;43mjax\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcore\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43meval_jaxpr\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 619\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mjaxpr\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mconsts\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43min_dynamic_flat\u001b[39;49m\n\u001b[32m 620\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 621\u001b[39m self_out_dynamic_struct = _unflatten(\u001b[38;5;28mself\u001b[39m.out_dynamic_struct)\n\u001b[32m 622\u001b[39m self_out_static = _unflatten(\u001b[38;5;28mself\u001b[39m.out_static)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:796\u001b[39m, in \u001b[36meval_jaxpr\u001b[39m\u001b[34m(jaxpr, consts, propagate_source_info, *args)\u001b[39m\n\u001b[32m 793\u001b[39m traceback = eqn.source_info.traceback \u001b[38;5;28;01mif\u001b[39;00m propagate_source_info \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m 794\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m (source_info_util.user_context(traceback, name_stack=name_stack),\n\u001b[32m 795\u001b[39m eqn.ctx.manager):\n\u001b[32m--> \u001b[39m\u001b[32m796\u001b[39m ans = \u001b[30;43meqn\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mprimitive\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mbind\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mmap\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mread\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43meqn\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43minvars\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mbind_params\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 797\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m eqn.primitive.multiple_results:\n\u001b[32m 798\u001b[39m foreach(write, eqn.outvars, ans)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:705\u001b[39m, in \u001b[36mPrimitive.bind\u001b[39m\u001b[34m(self, *args, **params)\u001b[39m\n\u001b[32m 703\u001b[39m trace_ctx.set_trace(eval_trace)\n\u001b[32m 704\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m705\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mbind_with_trace\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mprev_trace\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mavals\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 706\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m 707\u001b[39m trace_ctx.set_trace(prev_trace)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\custom_derivatives.py:1023\u001b[39m, in \u001b[36mCustomVJPCallPrimitive.bind_with_trace\u001b[39m\u001b[34m(self, trace, args, avals, params)\u001b[39m\n\u001b[32m 1021\u001b[39m params = \u001b[38;5;28mdict\u001b[39m(params)\n\u001b[32m 1022\u001b[39m fun, fwd, bwd = params.pop(\u001b[33m'\u001b[39m\u001b[33msubfuns\u001b[39m\u001b[33m'\u001b[39m)\n\u001b[32m-> \u001b[39m\u001b[32m1023\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mtrace\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mprocess_custom_vjp_call\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mfun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mfwd\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mbwd\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:1310\u001b[39m, in \u001b[36mEvalTrace.process_custom_vjp_call\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m 1308\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mprocess_custom_vjp_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, fun, fwd, bwd, tracers, /, **_):\n\u001b[32m 1309\u001b[39m \u001b[38;5;28;01mdel\u001b[39;00m primitive, fwd, bwd, _ \u001b[38;5;66;03m# Unused.\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m1310\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mfun\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcall_wrapped\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mtracers\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\linear_util.py:212\u001b[39m, in \u001b[36mWrappedFun.call_wrapped\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 210\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mcall_wrapped\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args, **kwargs):\n\u001b[32m 211\u001b[39m \u001b[38;5;250m \u001b[39m\u001b[33;03m\"\"\"Calls the transformed function\"\"\"\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m212\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mf_transformed\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:325\u001b[39m, in \u001b[36mjaxpr_as_fun\u001b[39m\u001b[34m(closed_jaxpr, *args)\u001b[39m\n\u001b[32m 316\u001b[39m \u001b[38;5;129m@curry\u001b[39m\n\u001b[32m 317\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mjaxpr_as_fun\u001b[39m(closed_jaxpr: ClosedJaxpr, *args):\n\u001b[32m 318\u001b[39m \u001b[38;5;66;03m# TODO(dougalm): remove this hack when we add contexts to jaxpr.\u001b[39;00m\n\u001b[32m (...)\u001b[39m\u001b[32m 322\u001b[39m \u001b[38;5;66;03m# context modifications. In the meantime, disabling the checks when we\u001b[39;00m\n\u001b[32m 323\u001b[39m \u001b[38;5;66;03m# round-trip prevents those ops producing spurious errors.\u001b[39;00m\n\u001b[32m 324\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m config.debug_nans(\u001b[38;5;28;01mFalse\u001b[39;00m):\n\u001b[32m--> \u001b[39m\u001b[32m325\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43meval_jaxpr\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mclosed_jaxpr\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mjaxpr\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mclosed_jaxpr\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mconsts\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:796\u001b[39m, in \u001b[36meval_jaxpr\u001b[39m\u001b[34m(jaxpr, consts, propagate_source_info, *args)\u001b[39m\n\u001b[32m 793\u001b[39m traceback = eqn.source_info.traceback \u001b[38;5;28;01mif\u001b[39;00m propagate_source_info \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m 794\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m (source_info_util.user_context(traceback, name_stack=name_stack),\n\u001b[32m 795\u001b[39m eqn.ctx.manager):\n\u001b[32m--> \u001b[39m\u001b[32m796\u001b[39m ans = \u001b[30;43meqn\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mprimitive\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mbind\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mmap\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mread\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43meqn\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43minvars\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mbind_params\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 797\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m eqn.primitive.multiple_results:\n\u001b[32m 798\u001b[39m foreach(write, eqn.outvars, ans)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:705\u001b[39m, in \u001b[36mPrimitive.bind\u001b[39m\u001b[34m(self, *args, **params)\u001b[39m\n\u001b[32m 703\u001b[39m trace_ctx.set_trace(eval_trace)\n\u001b[32m 704\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m705\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mbind_with_trace\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mprev_trace\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mavals\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 706\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m 707\u001b[39m trace_ctx.set_trace(prev_trace)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:713\u001b[39m, in \u001b[36mPrimitive.bind_with_trace\u001b[39m\u001b[34m(self, trace, args, avals, params)\u001b[39m\n\u001b[32m 711\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m set_current_trace(trace):\n\u001b[32m 712\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m.to_lojax(*args, **params) \u001b[38;5;66;03m# pyrefly: ignore[not-callable]\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m713\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mtrace\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mprocess_primitive\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\core.py:1294\u001b[39m, in \u001b[36mEvalTrace.process_primitive\u001b[39m\u001b[34m(self, primitive, args, params)\u001b[39m\n\u001b[32m 1292\u001b[39m args = \u001b[38;5;28mmap\u001b[39m(full_lower, args)\n\u001b[32m 1293\u001b[39m check_eval_args(args)\n\u001b[32m-> \u001b[39m\u001b[32m1294\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mprimitive\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mimpl\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\dispatch.py:88\u001b[39m, in \u001b[36mapply_primitive\u001b[39m\u001b[34m(prim, *args, **params)\u001b[39m\n\u001b[32m 86\u001b[39m prev = config.disable_jit.swap_local(\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[32m 87\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m---> \u001b[39m\u001b[32m88\u001b[39m outs = \u001b[30;43mfun\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 89\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m 90\u001b[39m config.disable_jit.set_local(prev)\n", + " \u001b[31m[... skipping hidden 1 frame]\u001b[39m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\pjit.py:259\u001b[39m, in \u001b[36m_cpp_pjit..cache_miss\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m 255\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mre-tracing function \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mjit_info.fun_sourceinfo\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m for \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 256\u001b[39m \u001b[33m\"\u001b[39m\u001b[33m`jit`, but \u001b[39m\u001b[33m'\u001b[39m\u001b[33mno_tracing\u001b[39m\u001b[33m'\u001b[39m\u001b[33m is set\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 257\u001b[39m p, args_flat = _infer_params(fun, jit_info, args, kwargs)\n\u001b[32m 258\u001b[39m (outs, out_flat, out_tree, args_flat, jaxpr,\n\u001b[32m--> \u001b[39m\u001b[32m259\u001b[39m executable, pgle_profiler, const_args) = \u001b[30;43m_run_python_pjit\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 260\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mp\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs_flat\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mfun\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 262\u001b[39m maybe_fastpath_data = _get_fastpath_data(\n\u001b[32m 263\u001b[39m executable, out_tree, args_flat, out_flat, jaxpr.effects, jaxpr.consts,\n\u001b[32m 264\u001b[39m pgle_profiler, const_args)\n\u001b[32m 266\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m outs, maybe_fastpath_data, _need_to_rebuild_with_fdo(pgle_profiler)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\pjit.py:139\u001b[39m, in \u001b[36m_run_python_pjit\u001b[39m\u001b[34m(p, args_flat, fun, args, kwargs)\u001b[39m\n\u001b[32m 137\u001b[39m args_flat = \u001b[38;5;28mmap\u001b[39m(core.full_lower, args_flat)\n\u001b[32m 138\u001b[39m core.check_eval_args(args_flat)\n\u001b[32m--> \u001b[39m\u001b[32m139\u001b[39m out_flat, compiled, profiler, const_args = \u001b[30;43m_pjit_call_impl_python\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 140\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs_flat\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mp\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mparams\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 141\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m 142\u001b[39m out_flat = jit_p.bind(*args_flat, **p.params)\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\pjit.py:1175\u001b[39m, in \u001b[36m_pjit_call_impl_python\u001b[39m\u001b[34m(jaxpr, in_shardings, out_shardings, in_layouts, out_layouts, donated_invars, ctx_mesh, name, keep_unused, inline, compiler_options_kvs, *args)\u001b[39m\n\u001b[32m 1164\u001b[39m arg_types = \u001b[38;5;28mmap\u001b[39m(convert_to_metaty, args)\n\u001b[32m 1165\u001b[39m computation = _resolve_and_lower(\n\u001b[32m 1166\u001b[39m arg_types, jaxpr=jaxpr, in_shardings=in_shardings,\n\u001b[32m 1167\u001b[39m out_shardings=out_shardings, in_layouts=in_layouts,\n\u001b[32m (...)\u001b[39m\u001b[32m 1173\u001b[39m compiler_options_kvs=compiler_options_kvs,\n\u001b[32m 1174\u001b[39m )\n\u001b[32m-> \u001b[39m\u001b[32m1175\u001b[39m compiled = \u001b[30;43mcomputation\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcompile\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1176\u001b[39m sharded_const_args = compiled.shard_const_args(computation.const_args)\n\u001b[32m 1178\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m config.distributed_debug.value:\n\u001b[32m 1179\u001b[39m \u001b[38;5;66;03m# Defensively only perform fingerprint logic if debug logging is enabled\u001b[39;00m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\interpreters\\pxla.py:1239\u001b[39m, in \u001b[36mMeshComputation.compile\u001b[39m\u001b[34m(self, compiler_options, device_assignment)\u001b[39m\n\u001b[32m 1236\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m use_cache \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m._executable \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m 1237\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._executable\n\u001b[32m-> \u001b[39m\u001b[32m1239\u001b[39m executable = \u001b[30;43mUnloadedMeshExecutable\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mfrom_hlo\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 1240\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43m_name\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43m_hlo\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcompile_args\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1241\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcompiler_options_kvs\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mcompiler_options_kvs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1242\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mdevice_list\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mcompilation_device_list\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1243\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m use_cache:\n\u001b[32m 1244\u001b[39m \u001b[38;5;28mself\u001b[39m._executable = executable\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\interpreters\\pxla.py:1733\u001b[39m, in \u001b[36mUnloadedMeshExecutable.from_hlo\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m 1729\u001b[39m allow_prop_to_inputs, allow_prop_to_outputs = get_prop_to_input_output(\n\u001b[32m 1730\u001b[39m in_shardings, out_shardings, \u001b[38;5;28mlen\u001b[39m(ordered_effects))\n\u001b[32m 1732\u001b[39m util.test_event(\u001b[33m\"\u001b[39m\u001b[33mpxla_cached_compilation\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m-> \u001b[39m\u001b[32m1733\u001b[39m xla_executable = \u001b[30;43m_cached_compilation\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 1734\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mhlo\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mname\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mtuple_args\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mallow_prop_to_inputs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mallow_prop_to_outputs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1735\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mtuple\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mhost_callbacks\u001b[39;49m\u001b[30;43m)\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mbackend\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mdevice_list\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcompiler_options_kvs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1736\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mpgle_profiler\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1738\u001b[39m in_shardings = _maybe_get_and_check_in_shardings(\n\u001b[32m 1739\u001b[39m xla_executable, in_shardings, device_list, global_in_avals,\n\u001b[32m 1740\u001b[39m \u001b[38;5;28mlen\u001b[39m(ordered_effects))\n\u001b[32m 1741\u001b[39m out_shardings = _maybe_get_and_check_out_shardings(\n\u001b[32m 1742\u001b[39m xla_executable, out_shardings, device_list, global_out_avals,\n\u001b[32m 1743\u001b[39m \u001b[38;5;28mlen\u001b[39m(ordered_effects))\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\interpreters\\pxla.py:1516\u001b[39m, in \u001b[36m_cached_compilation\u001b[39m\u001b[34m(computation, name, tuple_args, allow_prop_to_inputs, allow_prop_to_outputs, host_callbacks, backend, da, compiler_options_kvs, pgle_profiler)\u001b[39m\n\u001b[32m 1509\u001b[39m compile_options = create_compile_options(\n\u001b[32m 1510\u001b[39m computation, tuple_args, allow_prop_to_inputs, allow_prop_to_outputs,\n\u001b[32m 1511\u001b[39m backend, dev, compiler_options)\n\u001b[32m 1513\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m dispatch.log_elapsed_time(\n\u001b[32m 1514\u001b[39m \u001b[33m\"\u001b[39m\u001b[33mFinished XLA compilation of \u001b[39m\u001b[38;5;132;01m{fun_name}\u001b[39;00m\u001b[33m in \u001b[39m\u001b[38;5;132;01m{elapsed_time:.9f}\u001b[39;00m\u001b[33m sec\u001b[39m\u001b[33m\"\u001b[39m,\n\u001b[32m 1515\u001b[39m fun_name=name, event=dispatch.BACKEND_COMPILE_EVENT):\n\u001b[32m-> \u001b[39m\u001b[32m1516\u001b[39m xla_executable = \u001b[30;43mcompiler\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcompile_or_get_cached\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 1517\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mbackend\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcomputation\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mdev\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcompile_options\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mhost_callbacks\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 1518\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mda\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mpgle_profiler\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 1519\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m xla_executable\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\compiler.py:469\u001b[39m, in \u001b[36mcompile_or_get_cached\u001b[39m\u001b[34m(backend, computation, devices, compile_options, host_callbacks, executable_devices, pgle_profiler)\u001b[39m\n\u001b[32m 467\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m 468\u001b[39m log_persistent_cache_miss(module_name, cache_key)\n\u001b[32m--> \u001b[39m\u001b[32m469\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43m_compile_and_write_cache\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 470\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mbackend\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 471\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcomputation\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 472\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mexecutable_devices\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 473\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcompile_options\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 474\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mhost_callbacks\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 475\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mmodule_name\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 476\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcache_key\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 477\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\compiler.py:737\u001b[39m, in \u001b[36m_compile_and_write_cache\u001b[39m\u001b[34m(backend, computation, executable_devices, compile_options, host_callbacks, module_name, cache_key)\u001b[39m\n\u001b[32m 727\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_compile_and_write_cache\u001b[39m(\n\u001b[32m 728\u001b[39m backend: xc.Client,\n\u001b[32m 729\u001b[39m computation: ir.Module,\n\u001b[32m (...)\u001b[39m\u001b[32m 734\u001b[39m cache_key: \u001b[38;5;28mstr\u001b[39m,\n\u001b[32m 735\u001b[39m ) -> xc.LoadedExecutable:\n\u001b[32m 736\u001b[39m start_time = time.monotonic()\n\u001b[32m--> \u001b[39m\u001b[32m737\u001b[39m executable = \u001b[30;43mbackend_compile_and_load\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 738\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mbackend\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcomputation\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mexecutable_devices\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mcompile_options\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mhost_callbacks\u001b[39;49m\n\u001b[32m 739\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 740\u001b[39m compile_time = time.monotonic() - start_time\n\u001b[32m 741\u001b[39m _cache_write(\n\u001b[32m 742\u001b[39m cache_key, compile_time, module_name, backend, executable, host_callbacks\n\u001b[32m 743\u001b[39m )\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\profiler.py:420\u001b[39m, in \u001b[36mannotate_function..wrapper\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m 417\u001b[39m \u001b[38;5;129m@wraps\u001b[39m(func)\n\u001b[32m 418\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mwrapper\u001b[39m(*args, **kwargs):\n\u001b[32m 419\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m TraceAnnotation(name, **decorator_kwargs):\n\u001b[32m--> \u001b[39m\u001b[32m420\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mfunc\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n", + "\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\grace\\anaconda3\\envs\\linx\\Lib\\site-packages\\jax\\_src\\compiler.py:344\u001b[39m, in \u001b[36mbackend_compile_and_load\u001b[39m\u001b[34m(backend, module, executable_devices, options, host_callbacks)\u001b[39m\n\u001b[32m 342\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m 343\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m host_callbacks:\n\u001b[32m--> \u001b[39m\u001b[32m344\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mbackend\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mcompile_and_load\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m 345\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mmodule\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 346\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mexecutable_devices\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mexecutable_devices\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 347\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mcompile_options\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43moptions\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 348\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43mhost_callbacks\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mhost_callbacks\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m 349\u001b[39m \u001b[30;43m \u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m 350\u001b[39m \u001b[38;5;66;03m# Some backends don't have `host_callbacks` option yet\u001b[39;00m\n\u001b[32m 351\u001b[39m \u001b[38;5;66;03m# TODO(sharadmv): remove this fallback when all backends allow `compile`\u001b[39;00m\n\u001b[32m 352\u001b[39m \u001b[38;5;66;03m# to take in `host_callbacks`\u001b[39;00m\n\u001b[32m 353\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m backend.compile_and_load(\n\u001b[32m 354\u001b[39m module,\n\u001b[32m 355\u001b[39m executable_devices=executable_devices,\n\u001b[32m 356\u001b[39m compile_options=options,\n\u001b[32m 357\u001b[39m )\n", + "\u001b[31mKeyboardInterrupt\u001b[39m: " + ] + } + ], + "source": [ + "%load_ext autoreload\n", + "%autoreload\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import jax.numpy as jnp\n", + "import jax\n", + "jax.config.update(\"jax_disable_jit\", True)\n", + "from jax import jit, vmap\n", + "import sys\n", + "\n", + "sys.path.append(\"../\")\n", + "import linx.const as const \n", + "from linx.nuclear import NuclearRates\n", + "from linx.background import BackgroundModel\n", + "from linx.abundances import AbundanceModel\n", + "\n", + "thermo_model_DNeff = BackgroundModel(max_steps=1)\n", + "\n", + "(\n", + " t_vec_ref, a_vec_ref, rho_g_vec, rho_nu_vec, rho_NP_vec, P_NP_vec, Neff_vec \n", + ") = thermo_model_DNeff(0.)\n", + "\n", + "network = 'key_PRIMAT_2023'\n", + "# network = 'key_PRIMAT_2018'\n", + "# network = 'key_PArthENoPE'\n", + "# network = 'key_YOF'\n", + "abundance_model = AbundanceModel(NuclearRates(nuclear_net=network))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59e88d0d", + "metadata": {}, + "outputs": [], + "source": [ + "import linx.thermo as thermo\n", + "import linx.background as background\n", + "import linx.const as const\n", + "\n", + "Y = background.get_Y_from_abundances(abundance_model, thermo_model, 0.0)\n", + "t = np.arange(1, 10, 1)\n", + "diffeq = background.dY(t, Y[1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c9f3cd1a", + "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { - "display_name": "", - "name": "" + "display_name": "linx", + "language": "python", + "name": "python3" }, "language_info": { - "name": "" + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.15" } }, "nbformat": 4, diff --git a/linx/background.py b/linx/background.py index 4e07c21..9994fb5 100644 --- a/linx/background.py +++ b/linx/background.py @@ -133,11 +133,15 @@ def __call__( T_EM_init = T_start T_nu_init = T_EM_init + #grace, ask cara about this. + mu_nue_init = 0. + mu_numt_init = 0. + rho_extra_init = (7/8) * (4/11)**(4/3) * ( thermo.rho_massless_BE(T_EM_init, 0., 2) ) * Delt_Neff_init - Y0 = (lna_init, T_EM_init, T_nu_init) + Y0 = (lna_init, T_EM_init, T_nu_init, mu_nue_init, mu_numt_init) # use parametric form to estimate correct start # time given T_start, assuming T ~ t^(-1/2) and @@ -226,24 +230,28 @@ def dY(self, t, Y, args): ------- The time derivative of the quantities specified in `Y`. """ + - lna, T_g, T_nu = Y + lna, T_g, T_nu, mu_nue, mu_numt = Y + def printer(x): + print(x) + jax.debug.callback(printer,mu_nue) lna_init, rho_extra_init, me = args rho_EM = thermo.rho_EM_std(T_g, me=me, LO=self.LO, NLO=self.NLO) rho_plus_p_EM = thermo.rho_plus_p_EM_std(T_g, me=me, LO=self.LO, NLO=self.NLO) drho_EM_dT_g = thermo.drho_EM_dT_g_std(T_g, me=me, LO=self.LO, NLO=self.NLO) - rho_nu = 3*thermo.rho_nue_std(T_nu) + rho_nu = 3*thermo.rho_nue_std(T_nu, mu_nue) rho_plus_p_nu = (4/3) * rho_nu - drho_nu_dT_nu = 3*thermo.drho_nue_dT_nue_std(T_nu) + drho_nu_dT_nu = 3*thermo.drho_nue_dT_nue_std(T_nu, mu_nue) rho_extra = rho_extra_init * jnp.exp(lna_init)**4 / jnp.exp(lna)**4 H = thermo.Hubble(rho_EM + rho_nu + rho_extra) C_rho_nue, C_rho_numu, C_n_nue, C_n_numu = thermo.collision_terms_std( - T_g, T_nu, T_nu, me=me, decoupled=self.decoupled, use_FD=self.use_FD, collision_me=self.collision_me + T_g, T_nu, T_nu, mu_nue, mu_numt, decoupled=self.decoupled, use_FD=self.use_FD, collision_me=self.collision_me ) drho_EM_dt = -3 * H * rho_plus_p_EM - C_rho_nue - 2*C_rho_numu @@ -252,16 +260,19 @@ def dY(self, t, Y, args): dT_g_dt = drho_EM_dt / drho_EM_dT_g dT_nu_dt = drho_nu_dt / drho_nu_dT_nu - #P_e = thermo.rho_nue_std(T_g, me=me) / 3 - #dndt_e = C_n_nue(T_g, T_nu, T_nu, me=me, decoupled=self.decoupled, use_FD=self.use_FD, collision_me=self.collision_me) + P_e = thermo.rho_nue_std(T_g, mu_nue) / 3 + - #dmue_dt = -(-3.* H *( (thermo.rho_nue_std+P_e) * thermo.dn_nue_dT_nue_std - thermo.n_nue_std * thermo.drho_nue_dT_nue_std) + thermo.dn_nue_dT_nue_std * thermo.drho_nue_dT_nue_std - thermo.drho_nue_dT_nue_std * dndt_e )/(thermo.dn_nue_dmu_nue_std * thermo.drho_nue_dT_nue_std - thermo.dn_nue_dT_nue_std * thermo.drho_nue_dmu_nue_std) + # dmue_dt = -(-3.* H *( (thermo.rho_nue_std(T_nu)+P_e) * thermo.dn_nue_dT_nue_std(T_nu, mu_nue) - thermo.n_nue_std(T_nu) * thermo.drho_nue_dT_nue_std(T_nu, mu_nue)) + thermo.dn_nue_dT_nue_std(T_nu, mu_nue) * thermo.drho_nue_dT_nue_std(T_nu, mu_nue) - thermo.drho_nue_dT_nue_std(T_nu, mu_nue) * C_n_nue)/(thermo.dn_nue_dmu_nue_std(T_nu, mu_nue) * thermo.drho_nue_dT_nue_std(T_nu, mu_nue) - thermo.dn_nue_dT_nue_std(T_nu, mu_nue) * thermo.drho_nue_dmu_nue_std(T_nu, mu_nue)) - #P_mu = thermo.rho_numt_std(T_nu, me=me) / 3 - #dndt_mu = C_n_numu(T_g, T_nu, T_nu, me=me, decoupled=self.decoupled, use_FD=self.use_FD, collision_me=self.collision_me) + P_mu = thermo.rho_numt_std(T_nu, mu_numt=mu_numt) / 3 + + + # dmunu_dt = -(-3.*H*( (thermo.rho_numt_std(T_nu)+P_mu)* thermo.dn_numt_dT_numt_std(T_nu, mu_numt) - thermo.n_numt_std(T_nu) * thermo.drho_numt_dT_numt_std(T_nu, mu_numt)) + thermo.dn_numt_dT_numt_std(T_nu, mu_numt) * thermo.drho_numt_dT_numt_std(T_nu, mu_numt) - thermo.drho_numt_dT_numt_std(T_nu, mu_numt) * C_n_numu )/(thermo.dn_numt_dmu_numt_std(T_nu, mu_numt) * thermo.drho_numt_dT_numt_std(T_nu, mu_numt) - thermo.dn_numt_dT_numt_std(T_nu, mu_numt) * thermo.drho_numt_dmu_numt_std(T_nu, mu_numt)) - #dmunu_dt = -(-3.*H*( (thermo.rho_numt_std+P_mu)* thermo.dn_numt_dT_numt_std - thermo.n_numt.std * thermo.drho_numt_dT_numt_std) + thermo.dn_numt_dT_numt_std * thermo.drho_numt_dT_numt_std - thermo.drho_numt_dT_numt_std * dndt_mu )/(thermo.dn_numt_dmu_numt_std * thermo.drho_numt_dT_numt_std - thermo.dn_numt_dT_numt_std * thermo.drho_numt_dmu_numt_std) + dmue_dt = -(3. * H * (thermo.rho_nue_std(T_nu, mu_nue) + P_e)) - return H, dT_g_dt, dT_nu_dt - #return H, dT_g_dt, dT_nu_dt, dmue_dt, dmunu_dt + dmunu_dt = -(3. * H * (thermo.rho_numt_std(T_nu, mu_nue) + P_e)) + + return H, dT_g_dt, dT_nu_dt, dmue_dt, dmunu_dt diff --git a/linx/chemical potential testing.ipynb b/linx/chemical potential testing.ipynb index 2024abe..4274fda 100644 --- a/linx/chemical potential testing.ipynb +++ b/linx/chemical potential testing.ipynb @@ -46,7 +46,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "linx", "language": "python", "name": "python3" }, diff --git a/linx/linx chemical potential testing 4-30.ipynb b/linx/linx chemical potential testing 4-30.ipynb index 5607e5f..dcf5f75 100644 --- a/linx/linx chemical potential testing 4-30.ipynb +++ b/linx/linx chemical potential testing 4-30.ipynb @@ -67,7 +67,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 11, "id": "eef27bb3", "metadata": {}, "outputs": [ @@ -383,7 +383,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "linx", "language": "python", "name": "python3" },