Theano errors when using complex python functions inside PyMC3 models

Questions
1

I have an MCMC problem that requires calls to a python function that employs a zero-finder internally. It’s throwing errors that I’m not sure how to diagnose (I’m new to Theano)… any pointers greatly appreciated!

import pymc3 as pm
import cbsyst as cb  # the package I'm calling - for calculating seawater carbon chemistry.

# example problem: what DIC is consistent with an observed pCO2 and TA
obs_pCO2 = 433.5
TA = 2400

with pm.Model() as m:
    DIC = pm.Normal('DIC', 2000, 200)
    
    mu = cb.Csys(DIC=DIC, TA=TA).pCO2
    sigma = pm.Exponential('sigma', 1)
    
    pCO2 = pm.Normal('pCO2', mu, sigma, observed=obs_pCO2)
    
    pm.sample()

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-18-49e7ddbb65c4> in <module>
      3     DIC = pm.Normal('DIC', 2000, 200)
      4 
----> 5     mu = cb.Csys(DIC=DIC, TA=TA).pCO2
      6     sigma = pm.Exponential('sigma', 1)
      7 

~/GitHub/cbsyst/cbsyst/cbsyst.py in Csys(pHtot, DIC, CO2, HCO3, CO3, TA, fCO2, pCO2, BT, Ca, Mg, T_in, S_in, P_in, T_out, S_out, P_out, TP, TSi, pHsws, pHfree, pHNBS, Ks, pdict, unit)
    224                              fCO2=ps.fCO2, pCO2=ps.pCO2,
    225                              T_in=ps.T_in, BT=ps.BT, TP=ps.TP, TSi=ps.TSi,
--> 226                              TS=ps.TS, TF=ps.TF, Ks=ps.Ks))
    227 
    228     # calculate pHs on all scales, if not done before.

~/GitHub/cbsyst/cbsyst/carbon_fns.py in calc_C_species(pHtot, DIC, CO2, HCO3, CO3, TA, fCO2, pCO2, T_in, BT, TP, TSi, TS, TF, Ks)
    607                        TS=TS,
    608                        TF=TF,
--> 609                        Ks=Ks)
    610         H = ch(pHtot)
    611 

~/GitHub/cbsyst/cbsyst/carbon_fns.py in TA_DIC(TA, DIC, BT, TP, TSi, TS, TF, Ks)
    342 
    343         while any(abs(deltapH) > 1):
--> 344             deltapH[abs(deltapH) > 1] /= 2
    345 
    346         pHx += deltapH

~/miniconda3/envs/py3/lib/python3.7/site-packages/theano/tensor/var.py in __getitem__(self, args)
    556 
    557         if advanced_boolean:
--> 558             return theano.tensor.subtensor.advanced_boolean_subtensor(self, *args)
    559         elif advanced:
    560             if (axis is not None and

~/miniconda3/envs/py3/lib/python3.7/site-packages/theano/gof/op.py in __call__(self, *inputs, **kwargs)
    613         """
    614         return_list = kwargs.pop('return_list', False)
--> 615         node = self.make_node(*inputs, **kwargs)
    616 
    617         if config.compute_test_value != 'off':

~/miniconda3/envs/py3/lib/python3.7/site-packages/theano/tensor/subtensor.py in make_node(self, x, *index)
   2177 
   2178         index = tuple(map(as_index_variable, index))
-> 2179         bcast = adv_index_broadcastable_pattern(x, index)
   2180         return gof.Apply(self,
   2181                          (x,) + index,

~/miniconda3/envs/py3/lib/python3.7/site-packages/theano/tensor/subtensor.py in adv_index_broadcastable_pattern(a, idx)
   2114     # 2 - True = 1; 2 - False = 2
   2115     fakeshape = [2 - bc for bc in a.broadcastable]
-> 2116     retshape = np.empty(fakeshape)[newidx].shape
   2117     return tuple([dim == 1 for dim in retshape])
   2118 

IndexError: boolean index did not match indexed array along dimension 0; dimension is 1 but corresponding boolean dimension is 2

The error is definitely coming from inside the zero-finder part of the function. I’m open to modifying the zero-finder to work with Theano… but is this even possible? At this stage I’m just looking for some friendly pointers before spending a whole lot of time trying going down blind alleys!