Switching to pymc.math.softmax implementation; adding axis=1 argument to the softmax function for rethinking_2 chapter 11 notebook.

Author: DamonCharlesRoberts

Rethinking_2/Chp_11.ipynb

@@ -4719,7 +4719,7 @@
     "    s2 = 0.0 + b * income[2]  # pivoting the intercept for the third category\n",
     "    s = pm.math.stack([s0, s1, s2])\n",
     "\n",
-    "    p_ = at.nnet.softmax(s)\n",
+    "    p_ = pm.math.softmax(s, axis=1)\n",
     "    career_obs = pm.Categorical(\"career\", p=p_, observed=career)\n",
     "\n",
     "    trace_11_13 = pm.sample(tune=2000, target_accept=0.99, random_seed=RANDOM_SEED)\n",
@@ -5061,7 +5061,7 @@
     "    s2 = np.zeros(N)  # pivot\n",
     "    s = pm.math.stack([s0, s1, s2]).T\n",
     "\n",
-    "    p_ = at.nnet.softmax(s)\n",
+    "    p_ = pm.math.softmax(s, axis=1)\n",
     "    career_obs = pm.Categorical(\"career\", p=p_, observed=career)\n",
     "\n",
     "    trace_11_14 = pm.sample(1000, tune=2000, target_accept=0.9, random_seed=RANDOM_SEED)\n",