types changing in while_loop(), need help!

[derekpowell]

I'm fumbling my way through implementing some custom jvp for the regularized incomplete beta function by attempting to port over some code from the Stan project (this and this). The functions use a while loop that I am trying to implement with jax.lax.while_loop(). My implementations seem to work in some contexts but not others. When I try to run this inside a numpyro MCMC sampling context, I get the following error:

body_fun output and input must have identical types, got
('ShapedArray(float32[11])', 'DIFFERENT ShapedArray(float32[11]) vs. ShapedArray(float32[])', 'DIFFERENT ShapedArray(float32[11]) vs. ShapedArray(float32[])', 'ShapedArray(float32[])', 'ShapedArray(int32[], weak_type=True)', 'ShapedArray(float32[])', 'ShapedArray(float32[])', 'ShapedArray(float32[])', 'ShapedArray(float32[11])').

I gather I am somehow transforming the types inside my body_fun (_betainc_dda_while below) but I can't seem to figure out where. Hoping for some help debugging or even reading this error message, I can't see where the issue is arising.

here's my code:

import jax
import jax.nump as jnp
from jax import custom_jvp
from jax.scipy.special import gammaln
from jax.numpy import log1p, log, exp

@custom_jvp
def _betainc(a,b,x):
    return jax.scipy.special.betainc(a,b,x) # assign my own



def _betainc_dda_while(args):
    summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, digamma_a, x = args
    
    sum_numer = sum_numer + (digamma_ab - digamma_a) * summand
    sum_denom += summand
    summand = summand*(1 + (a_plus_b) / k) * (1 + k) / (1 + a_plus_1 / k)
    digamma_ab += 1./(a_plus_b + k)
    digamma_a += 1./(a_plus_1 + k)
    k += 1
    summand = summand * (x / k)

    args = (summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, digamma_a, x) 
    
    return args



## define partial derivative wrt a
def _betainc_dda(a_dot, primal_out, a, b, x):
    a = jnp.asarray(a)
    b = jnp.asarray(a)
    x = jnp.asarray(x)
    
    digamma_a = jax.scipy.special.digamma(a)
    digamma_ab = jax.scipy.special.digamma(a+b)
    
    threshold = 1e-10

    a_plus_b = a + b
    a_plus_1 = a + 1.
    
    digamma_a = digamma_a + 1./a
    
    prefactor = jnp.power(a_plus_1 / a_plus_b, 3)
    sum_numer = (digamma_ab - digamma_a) * prefactor
    sum_denom = prefactor
    summand = prefactor * x * a_plus_b / a_plus_1
    
    k = 1
    digamma_ab = digamma_ab + 1./a_plus_b
    digamma_a = digamma_a + 1./a_plus_1
    
    ### ----- 6/9/22, 4:03 PM something in the while loop changing the types?
    out = jax.lax.while_loop(
        lambda args: jnp.any(jnp.abs(args[0]) > 1e-10),
        _betainc_dda_while,
        (summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, digamma_a, x)
    )
    
    summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, digamma_a, x = out
    
    
    return _betainc(a, b, x) * (log(x) + sum_numer / sum_denom)*a_dot


def _betainc_ddb_while(args):
    summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, x = args
    
    sum_numer += digamma_ab * summand
    sum_denom += summand

    summand = summand*(1 + (a_plus_b) / k) * (1 + k) / (1 + a_plus_1 / k)
    digamma_ab += 1./(a_plus_b + k)
    k +=1
    summand = summand * x / k

    args = (summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, x)
    
    return args
        

## define partial derivative wrt b
def _betainc_ddb(b_dot, primal_out, a, b, x):
    
    digamma_b = jax.scipy.special.digamma(b)
    digamma_ab = jax.scipy.special.digamma(a+b)
    
    threshold = 1e-10
    
    a_plus_b = a + b
    a_plus_1 = a + 1.
    
    prefactor = jnp.power(a_plus_1 / a_plus_b, 3)
    
    sum_numer = digamma_ab * prefactor
    sum_denom = prefactor
    summand = prefactor * x * a_plus_b / a_plus_1
    
    k = 1
    digamma_ab = digamma_ab + 1./a_plus_b
    
    out = jax.lax.while_loop(
        lambda args: jnp.any(jnp.abs(args[0]) > 1e-10),
        _betainc_ddb_while,
        (summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, x)
    )
    
    summand, sum_numer, sum_denom, a_plus_b, k, a_plus_1, digamma_ab, x = out
        
    
    return _betainc(a, b, x) * (log(1 - x) - digamma_b + sum_numer / sum_denom)*b_dot




def betainc_gradx(g, primal_out, a, b, x):
    lbeta = gammaln(a) + gammaln(b) - gammaln(a + b)
    partial_x = exp((b - 1) * log1p(-x) +
                  (a - 1) * log(x) - lbeta)
    return partial_x * g

_betainc.defjvps(_betainc_dda, _betainc_ddb, betainc_gradx)

## these work
print(grad(_betainc, 0)(2., 1., .6)) 
print(grad(_betainc, 1)(2., 1., .6))
print(grad(_betainc, 2)(2., 1., .6))

Appreciate any help!

[soraros]

The error message says input sum_numer and sum_denom have the wrong shape: ShapedArray(float32[11]) vs. ShapedArray(float32[]). It's clearer if you reformat the error message a bit.

('ShapedArray(float32[11])',                                      # summand
 'DIFFERENT ShapedArray(float32[11]) vs. ShapedArray(float32[])', # sum_numer
 'DIFFERENT ShapedArray(float32[11]) vs. ShapedArray(float32[])', # sum_denom
 'ShapedArray(float32[])',                                        # a_plus_b
 'ShapedArray(int32[], weak_type=True)',                          # k
 'ShapedArray(float32[])',                                        # a_plus_1
 'ShapedArray(float32[])',                                        # digamma_ab
 'ShapedArray(float32[])',                                        # digamma_a
 'ShapedArray(float32[11])')                                      # x

Maybe you could also share the code that produces this error?

[derekpowell]

Ah that is very helpful! Than you, that was the first part of my question, how to read the error. So I gather I'm inadvertently doing something to the shapes of sum_numer and sum_denom?

Need to give a bit more context and show a bit more code to get to the error itself. I’m trying to model data from a survey where people were asked to report probabilities. The underlying cognitive model involves a Beta distribution, but the responses are often rounded. So the efforts to implement the Beta CDF relate to the rounding.

Here’s the numpyro data simulation and model I’m using for testing:

from jax.random import PRNGKey

X_raw = dist.Beta(.75*10, (1-.75)*10).sample(PRNGKey(10), (100,))
x_round = round(X_raw*10)

def f(mu, k):
    a = mu*k
    b = (1.-mu)*k
    
    responses = jnp.linspace(0,10, num=11)
    lower = jnp.clip((responses/10.) - .05, 0., 1.)
    upper = jnp.clip((responses/10.) + .05, 0., 1.)

    prob_resps = _betainc(a, b, upper) - _betainc(a, b, lower)
    
    return(prob_resps)


def mymodel_round(x=None):
    mu = numpyro.sample("mu", dist.Beta(1,1)) # noise parameter
    k = numpyro.sample("k", dist.HalfCauchy(10)) # noise parameter

    resp_probs = f(mu,k)
    
    with numpyro.plate("data", x.shape[0]):

        xhat = numpyro.sample("xhat", dist.Categorical(probs=resp_probs), obs=x)

And the error comes when I try to draw samples from the model in this code below. The trace says the specific error is occurring in _betainc_dda (presumably it will also occur in _betainc_ddb as well but just doesn't get there?)

kernel = NUTS(mymodel_round, target_accept_prob=.80)

mcmc = MCMC(kernel, 
               num_warmup=1_000, 
               num_samples=1_000, 
               num_chains=1)

mcmc.run(random.PRNGKey(0), x_round)

[soraros]

Haven't run your code, but I think the problem lies in the following line

prob_resps = _betainc(a, b, upper) - _betainc(a, b, lower)

where your upper and lower are of shape float32[11], while in the JVP definition, the assumption seems to be that they are scalars. You could either vmap _betainc over upper/lower, or use full_like in _betainc_dda_while (something like full_like(x, sum_numer).

[derekpowell]

Thank you, seems you are right! I added full_like in _betainc_dda as recommended and that's definitely doing something. The model still won't sample properly but now the functions no longer error.

I'm a bit confused about why that was needed though and how to test for this kind of thing in the future? Before and after this change, when I pass arrays straight into my _betainc_dda() function it handles them just fine. But when I set them up as custom jvp and call them with grad they don't. That is, this works:

_betainc_dda(jnp.array([1.,1.,1.]), None, jnp.array([3., 5, 1]), jnp.array([3., 5, 1]), jnp.array([.5,.6, .9]))

>  DeviceArray([-0.17469856, -0.09795534, -0.09482451], dtype=float32)

But this doesn't

grad(_betainc, 0)(jnp.array([2.,3]), jnp.array([2.,3]), jnp.array([.5,.7]))

>  TypeError: Gradient only defined for scalar-output functions. 

So I guess I must not be understanding something about custom jvp and grad?