Create Deterministic from pymc3 model

Questions
1

I am currently trying to have a distribution over a mutual information.
The goal is to ideally find the distributions that maximizes mutual information

I currently managed to do this using a multiple pymc3 models, but the sampling becomes quite expensive.
So I was wondering if there was a different approach compared to the following:

from sklearn.feature_selection import mutual_info_regression
import pymc3 as pm
import theano.tensor as tt
@theano.compile.ops.as_op(itypes=[tt.dscalar, tt.dscalar],
                          otypes=[tt.dscalar])
def test(mu, sigma):
    i = np.random.randint(0,100_000_000)
    with pm.Model() as model:
        A = pm.Uniform(f"A{i}", 0, 100)
        O = pm.Normal(f"O{i}", mu, sigma)

        out = pm.Deterministic(f"out{i}", f(A,O))

        trace = pm.sample(10_000)
        return mutual_info_regression(trace[f"A{i}"].reshape(-1,1), trace[f"out{i}"])[0]

with pm.Model() as model2:
    theta_1 = pm.Uniform("t1", 0, 100)
    theta_2 = pm.Uniform("t2", 0.001, 1_000)
    
    output = pm.Deterministic("output", test(theta_1,theta_2))
    
    trace = pm.sample(10_000)