Vmap, Scan, and numpyro.sample

tommcandrew · November 6, 2023, 2:53am

Hi All,
i need a bit of help adding in a numpyro.sample statement (or an equivalent stochastic statement) inside of jax.lax.scan.

The steps of the code that are giving me trouble are below.
The steps are

Vmap across a set of parameters
Inside the Vmap, there exists a scan
Inside the scan there is a statement: stochastic_states = numpyro.sample(f'stochastic_states_{t}', dist.Dirichlet(y) )

If i run the below code using jax.lax.scan a “UnexpectedTracerError” error is returned
If i run the below code replacing the jax.lax.scan with a scan from numpyro.contrib.control_flow then a shape mismatch error is returned

The code works/runs without the sample statement.

How do i include a stochastic component inside my scan?

Thanks for the help

        def step(carry,arr, M,ttls, beta_a1_beta2_a2):
            s,I,i,H,h,R = carry
            t,beta1,ky = arr

            beta2,a1,a2,alpha_removed_i,alpha_removed_h = beta_a1_beta2_a2

            a1  = beta1.reshape(num_locations,1)*(a1.reshape(num_locations,LAG))
            a2  = (a2.reshape(num_locations,LAG))

            contributions = ( ((i*a1).sum(1)).reshape(num_locations,1)  )
            l             = jnp.exp(-1*jnp.matmul(M,contributions)  )

            #--iniital proportions
            init_s   = s[:,0].reshape(-1,1)
            init_I   = I[:,0].reshape(-1,1)
            init_i   = i[:,0].reshape(-1,1)
            init_H   = H[:,0].reshape(-1,1)
            init_h   = h[:,0].reshape(-1,1)
            init_R   = R[:,0].reshape(-1,1)

            #--new proportions
            new_s     = (init_s*l).reshape(-1,1)

            new_i    = (init_s*(1.-l)).reshape(-1,1)
            new_I    = (init_I+new_i - alpha_removed_i*init_I).reshape(-1,1)
            
            new_h     = (a2*i).sum(1).reshape(num_locations ,1)
            new_H     = (init_H+new_h - alpha_removed_h*init_H).reshape(-1,1)

            new_R     = (init_R + alpha_removed_i*new_I + alpha_removed_h*new_H).reshape(-1,1)
            
            #--add stochastic component
            y = jnp.hstack([new_s,new_I,new_i, new_H, new_h, new_R])#.T
            
           #THIS IS THE PROBLEM
            stochastic_states   = numpyro.sample(f'stochastic_states_{t}', dist.Dirichlet(y) )
           #---
            x = jnp.delete(carry,-1,axis=2)
            y = stochastic_states.T

            z = jnp.zeros( (6,num_locations,LAG) )
            z = z.at[..., 1: ].set( x )
            z = z.at[...,0].set(y.reshape(6,num_locations))

            return z,y
        #----------------------------------------------------------------
        #--COLLECT STATES------------------------------------------------
        ttls = jnp.matmul(M,N).reshape(-1,1)
        def take_step(z,inits):
            s0,I0,i0,H0,h0,R0 = inits

            b1,beta2,a1,a2,alpha_removed_i,alpha_removed_h = z
            z=(beta2,a1,a2,alpha_removed_i,alpha_removed_h)

            states     = jnp.stack([s0,I0,i0,H0,h0,R0])
            rints      = jax.random.split( jax.random.PRNGKey(np.random.randint(0,1000)) ,time_units)
            def stoch_step(x,y):
                return step(x,y,M,ttls,z)
            _,scanned  = jax.lax.scan( stoch_step , init = states, xs = (jnp.arange(time_units),b1.T,rints))

            states = jnp.moveaxis(states,[2],[0])
            rslt = jnp.vstack( [ states, scanned])
            return rslt[-(time_units):,...]
        states=jax.vmap( lambda a,b: numpyro.handlers.seed( lambda z,inits: take_step(z,inits), jax.random.PRNGKey(1))( a,b )    )( (beta1,beta2.T,a1.T,a2.T,alpha_removed_i,alpha_removed_h), (s0,I0,i0,H0,h0,R0) )

fehiepsi · November 6, 2023, 10:36am

Hi @tommcandrew, could you make a small repro code?

tommcandrew · November 6, 2023, 1:46pm

Absolutely. i’ll make one today and give a fuller report

tommcandrew · November 6, 2023, 5:13pm

Hi @fehiepsi

i have built a minimal example but it looks like i bumped into another error.
The bug is in the line stochastic_states = numpyro.sample(f'stochastic_state_{t}', dist.Dirichlet(last_state) )

It looks like t is recognized as a tracer value and not a typical float so that i can rename this site iteratively.

If that line is commented out then the code runs.

Thanks again for the help,
tom

import numpy as np
import jax
from jax import jit
import jax.numpy as jnp

import numpyro
import numpyro.distributions as dist
from numpyro.contrib.control_flow import scan
import numpyro.handlers


if __name__ == "__main__":

    def example():
        def step(carry,arr):
            state = carry
            t = arr

            last_state = state[:,0]
            stochastic_states   = numpyro.sample(f'stochastic_state_{t}', dist.Dirichlet(last_state) )

            x = jnp.delete(carry,-1,axis=1)
            y = last_state #stochastic_states.T.reshape(5,1)

            z = jnp.zeros( (5,10) )
            z = z.at[:,1:].set( x )
            z = z.at[:,0].set(y.reshape(5,))
            
            return z,y
        #----------------------------------------------------------------
        #--COLLECT STATES------------------------------------------------
        z = jnp.zeros( (5,10) )

        #option one with error
        states = jax.lax.scan( step, init=z, xs=(jnp.arange(45)) )

        #option two with error
        states = scan( step, init=z, xs=(jnp.arange(45)) )
        

    from numpyro.infer import MCMC, NUTS,init_to_median
    mcmc = MCMC(
        NUTS(example
             , dense_mass=False, max_tree_depth=5,init_strategy = init_to_median(num_samples=100) ) #)init_to_sample() )
        , num_warmup  = 100
        , num_samples = 200
        , num_chains  = 1
        , thinning    = 2
    )
    mcmc.run(jax.random.PRNGKey(20200320))

fehiepsi · November 6, 2023, 5:44pm

Yes, t is a tracer. I think you don’t need to specify t in the name. See the example in scan doc.

tommcandrew · November 6, 2023, 6:37pm

@fehiepsi i thought the same, but when i remove t then a “repeated site name” error is returned.

tommcandrew · November 7, 2023, 1:35am

Update @fehiepsi

i have an updated example.
The code in example1 runs without error.
The code in example2 returns an error with an unexpected tracer.
The code in example3 returns no error
The code in example4 returns an error: AssertionError: all sites must have unique names but got stochastic_state duplicated

Example1 i call the scan without using vmap
Example2 uses vmap
Example3 places the scan inside of another function. This example was to exclude the possibility that placing scan inside of another function could cause an issue.
Example4 tries to replace the vmap with a simpler for loop.

Is there something special about vmap (or for loop in example4) that i need to take care of for this to work?

thanks for the help,
tom




#mcandrew
import numpy as np
import jax
from jax import jit
import jax.numpy as jnp

import numpyro
import numpyro.distributions as dist
from numpyro.contrib.control_flow import scan
import numpyro.handlers


if __name__ == "__main__":

    def example1():
        def step(carry,arr):
            state = carry
            t = arr

            last_state = state[:,0]
            stochastic_states   = numpyro.sample("stochastic",dist.Normal(jnp.ones(5,),1)) #numpyro.sample('stochastic_state', dist.Dirichlet(last_state) )

            x = jnp.delete(carry,-1,axis=1)
            y = stochastic_states.T.reshape(5,1)

            z = jnp.zeros( (5,10) )
            z = z.at[:,1:].set( x )
            z = z.at[:,0].set(y.reshape(5,))
            
            return z,y
        #----------------------------------------------------------------
        #--COLLECT STATES------------------------------------------------
        z = jnp.ones( (5,10) )

        #option two with error
        _,states = scan( step, init=z, xs=(jnp.arange(45)) )
        print(states.shape)
        
    def example2():
        def step(carry,arr):
            state = carry
            t = arr

            last_state = state[:,0]
            stochastic_states   = numpyro.sample('stochastic_state', dist.Dirichlet(last_state) )

            x = jnp.delete(carry,-1,axis=1)
            y = stochastic_states.T.reshape(5,1)

            z = jnp.zeros( (5,10) )
            z = z.at[:,1:].set( x )
            z = z.at[:,0].set(y.reshape(5,))
            
            return z,(y,x)
        #----------------------------------------------------------------
        #--COLLECT STATES------------------------------------------------
        z = jnp.ones( (5,10) )

        #option two with error
        def to_map(z):
            _,states = scan( step, init=z, xs=(jnp.arange(45)) )
            return states
        states = jax.vmap( to_map )( jnp.repeat( z[jnp.newaxis,...],3,axis=0 ) )

    def example3():
        def step(carry,arr):
            state = carry
            t = arr

            last_state = state[:,0]
            stochastic_states   = numpyro.sample('stochastic_state', dist.Dirichlet(last_state) )

            x = jnp.delete(carry,-1,axis=1)
            y = stochastic_states.T.reshape(5,1)

            z = jnp.zeros( (5,10) )
            z = z.at[:,1:].set( x )
            z = z.at[:,0].set(y.reshape(5,))
            
            return z,(y,x)
        #----------------------------------------------------------------
        #--COLLECT STATES------------------------------------------------
        z = jnp.ones( (5,10) )

        #option two with error
        def to_map(z):
            _,states = scan( step, init=z, xs=(jnp.arange(45)) )
            return states
        states =  to_map(z)

    def example4():
        def step(carry,arr,t):
            state = carry
            t = arr

            last_state = state[:,0]
            stochastic_states   = numpyro.sample('stochastic_state', dist.Dirichlet(last_state) )

            x = jnp.delete(carry,-1,axis=1)
            y = stochastic_states.T.reshape(5,1)

            z = jnp.zeros( (5,10) )
            z = z.at[:,1:].set( x )
            z = z.at[:,0].set(y.reshape(5,))
            
            return z,(y,x)
        #----------------------------------------------------------------
        #--COLLECT STATES------------------------------------------------
        z = jnp.ones( (5,10) )

        #option two with error
        def to_map(z,t):
            _,states = scan( lambda x,y :step(x,y,t), init=z, xs=(jnp.arange(45)) )
            return states
        
        for t,z in enumerate(jnp.repeat( z[jnp.newaxis,...],3,axis=0 )):
            states =  to_map(z, t)
            
    from numpyro.infer import MCMC, NUTS,init_to_median
    mcmc = MCMC(
        NUTS(example4
             , dense_mass=False, max_tree_depth=5,init_strategy = init_to_median(num_samples=100) ) #)init_to_sample() )
        , num_warmup  = 100
        , num_samples = 200
        , num_chains  = 1
        , thinning    = 2
    )
    mcmc.run(jax.random.PRNGKey(20200320))

fehiepsi · November 7, 2023, 2:07pm

re for loop: we need to specify different latent names in its body

re vmap: we don’t support vmap over non-closed numpyro programs. You might consider using plate in the scan body function here.

mbiron · July 28, 2026, 1:21am

Hi @fehiepsi – sorry for reviving this ~3 year old question, but your very last sentence is the only clue I’ve found online as to why the following doesn’t work

import jax
from jax import numpy as jnp

import numpyro
from numpyro import distributions as dist
from numpyro.infer import MCMC,HMC,NUTS

# define model
def postprocess(coin, z):
    x = z + jnp.where(coin, 1.0, -1.0)
    numpyro.deterministic('x', x)
    numpyro.sample('y', dist.Normal(x), obs=0.0)

def mixture_model():
    z = numpyro.sample('z', dist.Normal())
    coin = numpyro.sample('coin', dist.Bernoulli(0.5))
    if jnp.shape(coin) == ():
        postprocess(coin, z)
    else: # handle parallel enumeration
        jax.vmap(postprocess, in_axes=(0,None))(coin, z)

# sample
kernel = NUTS(mixture_model)
mcmc = MCMC(kernel, num_warmup=0, num_samples=1, progress_bar=False)
mcmc.run(jax.random.key(5))

Is postprocess a non-closed numpyro program? Because when I run this I get an UnexpectedTracerError (only including the last part of the stack trace for clarity)

 File "[...]/.venv/lib/python3.12/site-packages/numpyro/contrib/funsor/infer_util.py", line 222, in _enum_log_density
    log_prob = site["fn"].log_prob(value)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/distributions/util.py", line 799, in wrapper
    log_prob = log_prob_fn(self, *args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/numpyro/distributions/continuous.py", line 2471, in log_prob
    value_scaled = (value - self.loc) / self.scale
                    ~~~~~~^~~~~~~~~~
  File "[...]/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 1532, in op
    return getattr(self.aval, f"_{name}")(self, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 766, in _operator_rsub
    return ufuncs.subtract(cast(ArrayLike, other), self)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/.venv/lib/python3.12/site-packages/jax/_src/numpy/ufunc_api.py", line 182, in __call__
    return call(*args)
           ^^^^^^^^^^^
jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float32[] wrapped in a BatchTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.UnexpectedTracerError

In this toy example I can just remove the if/else block since jnp.where supports broadcasting. But in my actual application this is not an option. Is there an alternative way to rewrite this that would make it legal?