Something changed in `pytensor > 2.12.3` (and thus `pymc > 5.6.1`) that makes my `pytensor.gradient.grad` call get stuck - any ideas?

Questions v5
25

Here’s what I came up with:

#     # 7. Jacobian adjustment for observed data
    # Situation: y_c is an (n, 2) matrix of equations: [[eq[0, 0], eq[0, 1]], [eq[1, 0], eq[1, 1]], ...]
    # We ravel it to an (n * k, ) vector of equations: [eq[0, 0], eq[0, 1], [eq[1, 0], ...]
    # The vector ends up with entries sorted by equation, with entries alternating between the columns
    # Then we take the jacobian w.r.t the inputs. This will be (n * k, n, k)
    # Since the rows of equations are independent, this will have a sparsity structure that will allow us to undo the ravel in the next step.
    jacs = tg.jacobian(y_c.ravel(), y)
    
    # To unravel the jacobian, reshape to (n, 2, n, 2)
    # We need to merge away the 2nd batch dimension to get (n, k, k)
    # Because all the rows are independent, we can just use .sum() to accomplish this
    # (You can do jacs.eval() and confirm that all entries in the 2nd dimension are zero -- it contains the "cross terms" between the equations) 
    jacs = jacs.reshape((n, k, n, k)).sum(axis=2)
    
    # Now we're ready to get the logdets
    signs, logdets = pt.nlinalg.slogdet(jacs)
    
    _ = pm.Potential('pot_jcd_y_c_on_y',logdets, dims='oid')

To confirm this ends up with the correct answer, I added the first row of dfy as a separate variable and computed its jacobian alone. You can confirm that jacs[0] is equal to tg.jacobian(yc_1, y1) in the follow snippet:

n = cfg['n']
k = len(cfg['mu'])

with pm.Model(coords=coords, coords_mutable=coords_m) as mdl:
    y = pm.MutableData('y', dfy, dims=('oid', 'y_nm'))
    y0 = pm.MutableData('y1', dfy.loc[0, :])

    # 1. Define Observed Y marginal dists
    m_mu = pm.Normal('m_mu', mu=0.0, sigma=1.0, dims='m_nm')
    m_s = pm.InverseGamma('m_s', alpha=5.0, beta=4.0, dims='s_nm')
    m_d = pm.LogNormal.dist(mu=m_mu, sigma=m_s, shape=(cfg['n'], 2))

    m_d0 = pm.LogNormal.dist(mu=m_mu, sigma=m_s, shape=(1, 2))
    
    # 2. Evidence Observed Y against marginal PDFs
    _ = pm.Potential('pot_yhat', pm.logp(m_d, y), dims=('oid', 'yhat_nm'))

    # 3. Transform Observed Y to Uniform through marginal CDFs
    y_u = pm.Deterministic('y_u', pt.exp(pm.logcdf(m_d, y)))
    y_u1 = pm.Deterministic('y_u1', pt.exp(pm.logcdf(m_d0, y0)))
    

    # 4. Transform Uniformed Y to an Empirical Copula via MvN InvCDF
    y_c = pm.Deterministic('y_c', normal_icdf(y_u))
    y_c1 = pm.Deterministic('y_c1', normal_icdf(y_u1))

Aside: I think it’s a bug (or at least a missing feature) that you can’t ask for the gradient of a slice with respect to another slice when we know the computations are all independent. This is where the DisconnectedInput errors were coming from.

Edit: I made this gist to show what’s going on with all the ravels and reshapes, maybe it’s helpful? Jacobians of high-dimension (but independent) systems · GitHub

2 Likes