Thanks so much for the help. The marginalization of z to use nuts is super helpful!
Unfortunately regarding the subj index to get around the for loop – my variable m is a matrix of 14,000 x 2, not 50 x 2. In other words, each trial (out of 14,000 trials) has a different value for m[trial,1]. It’s not constant across all trials for each subject (although m[:,0] is a constant value, if that matters). So if I use the code as you adjusted it, I get a shape problem: “Input dimension mis-match. (input[0].shape[0] = 14000, input[1].shape[0] = 50)”. Do you know if there a workaround for this?