How indexing works in pymc

Questions v5
4

I’m on pymc version 4.1.3.

I think the only thing missing was import statements. If you run this, it should throw the broadcasting error I mentioned.

import pymc as pm
import arviz as az
import numpy as np
from scipy import stats

# synthetic data
# one predictor variable 
# representing time

x = np.arange(22)

# n is a 2x22 matrix representing
# the maximiums at each point in time

n = np.random.randint(100,size=(2,22))

# apply inverse logit transformation on 
# x twice with two different paramters

a0 = -4
b0 = 0.1

p0 = np.exp(a0 + b0*x) / (1 + np.exp(a0 + b0*x))

a1 = -2
b1 = 0.08

p1 = np.exp(a1 + b1*x) / (1 + np.exp(a1 + b1*x))

# generate outcome variable,
# two count variables from
# binomial distribution

k0 = stats.binom(n=n[0],p=p0).rvs()
k1 = stats.binom(n=n[1],p=p1).rvs()

# combine to get 2 x 22 outcome array

k = np.array([k0,k1])

cluster = np.array([0,1])

with pm.Model() as partial_pool:
    hyper_b_mu = pm.Normal('hbmu',0,1)
    hyper_b_std = pm.HalfNormal('hbstd',0.5)
    
    a = pm.Normal('a',0,1,shape=2)
    b = pm.Normal('b',hyper_b_mu,hyper_b_std,shape=2)

    p = pm.invlogit(a[cluster]+b[cluster]*x)
    
    y = pm.Binomial('y',n=n[cluster],p=p[cluster],observed=k[cluster])
    
    trace_partial_pool = pm.sample()

The short version of the error looks like:

ERROR (aesara.graph.opt): Optimization failure due to: constant_folding
ERROR (aesara.graph.opt): node: Assert{msg=Could not broadcast dimensions}(ScalarConstant{22}, TensorConstant{False})
ERROR (aesara.graph.opt): TRACEBACK:
ERROR (aesara.graph.opt): Traceback (most recent call last):
  File "C:\Users\Daniel Saunders\anaconda3\envs\pymc_env\lib\site-packages\aesara\graph\opt.py", line 1861, in process_node
    replacements = lopt.transform(fgraph, node)
  File "C:\Users\Daniel Saunders\anaconda3\envs\pymc_env\lib\site-packages\aesara\graph\opt.py", line 1066, in transform
    return self.fn(fgraph, node)
  File "C:\Users\Daniel Saunders\anaconda3\envs\pymc_env\lib\site-packages\aesara\tensor\basic_opt.py", line 2785, in constant_folding
    required = thunk()
  File "C:\Users\Daniel Saunders\anaconda3\envs\pymc_env\lib\site-packages\aesara\graph\op.py", line 559, in rval
    r = p(n, [x[0] for x in i], o, params)
  File "C:\Users\Daniel Saunders\anaconda3\envs\pymc_env\lib\site-packages\aesara\raise_op.py", line 96, in perform
    raise self.exc_type(self.msg)
AssertionError: Could not broadcast dimensions