Using jax.pmap inside jax.lax.scan for multicore computation of the loop

[JaySandesara]

Hello,

Apologies if this is a naive question, but I am trying to explore faster ways to compile and execute a jit-compiled function, that requires O(100k) repeated calls. In order to perform faster, multi-core computations I am trying to use jax.pmap for parallel computation of jax.lax.scan loop inside the jit-compiled function of interest. However, I have been having trouble understanding how exactly to transform the inputs and the loop to perform such a computation with jax.pmap. Here is a much simplified version of the code, I would appreciate any insights/feedback or even an alternative suggestion.

import os, sys, importlib

import numpy as np
import multiprocessing

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count={}".format(
    multiprocessing.cpu_count()
)

import jax
from jax import numpy as jnp
from functools import partial
from jax import pmap

jax.config.update('jax_platform_name', 'cpu')

platform = jax.lib.xla_bridge.get_backend().platform.casefold()
print("Platform: ", platform)

n_devices = jax.local_device_count()

# Random inputs - I have the flexibility to request as many cores as the rows of the inputs. 

arr_ratio_up_x = np.ones((n_devices, 341749))
arr_ratio_down_x = np.ones((n_devices, 341749))

arr_ratio_up_y = np.ones((n_devices, 34179))
arr_ratio_down_y = np.ones((n_devices, 34179))

ones_x = np.ones((arr_ratio_up_x.shape[1],))
ones_y = np.ones((arr_ratio_up_y.shape[1],))

tuple_init = (ones_x,ones_y)

combined_var_up = jnp.column_stack((arr_ratio_up_x, arr_ratio_up_y))
combined_var_down = jnp.column_stack((arr_ratio_down_x, arr_ratio_down_y))

len_x = arr_ratio_up_x.shape[1]
len_y = arr_ratio_up_y.shape[1]

tuple_init_alt = np.ones(combined_var_up.shape[1])

# Polynomial interpolation functions - Ignore for this example. Could be any function
@jax.jit
def poly_interp(tuple_input):

    alpha, varUp, varDown, varNom = tuple_input

    pow_up = jnp.divide(varUp,varNom)
    pow_down = jnp.divide(varDown,varNom)
    logHi = jnp.log(jnp.clip(pow_up,1e-7,None))
    logLo = jnp.log(jnp.clip(pow_down,1e-7,None))
    logHi = jnp.log(pow_up)
    logLo = jnp.log(pow_down)
    pow_up_log = jnp.multiply(pow_up,logHi)
    pow_down_log = -jnp.multiply(pow_down,logLo)
    pow_up_log2 = jnp.multiply(pow_up_log,logHi)
    pow_down_log2 = -jnp.multiply(pow_down_log,logLo)

    S0 = (pow_up+pow_down)/2.0
    A0 = (pow_up-pow_down)/2.0
    S1 = (pow_up_log+pow_down_log)/2.0
    A1 = (pow_up_log-pow_down_log)/2.0
    S2 = (pow_up_log2+pow_down_log2)/2.0
    A2 = (pow_up_log2-pow_down_log2)/2.0

    a1 = (15*A0-7*S1+A2)/8.0
    a2 = (-24+24*S0-9*A1+S2)/8.0
    a3 = (-5*A0+5*S1-A2)/4.0
    a4 = (12-12*S0+7*A1-S2)/4.0
    a5 = (3*A0-3*S1+A2)/8.0
    a6 = (-8+8*S0-5*A1+S2)/8.0

    return 1.0 + alpha * (a1 + alpha * ( a2 + alpha * ( a3 + alpha * ( a4 + alpha * ( a5 + alpha * a6 ) ) ) ) )

@jax.jit
def exp_interp(tuple_input):

    alpha, varUp, varDown, varNom = tuple_input

    return jnp.where(alpha>1.0, (varUp/varNom)**alpha, (varDown/varNom)**(-alpha))


# Toy function - ignore the details
@jax.jit
def sigma(mu, mu_0, mu_1, mu_2, ratio_x, ratio_y):

    return ratio_x.sum()+ratio_y.sum()

# Function of interest - simplified for this question
@jax.jit
def func_optimized_scan(param_array, combined_var_up, combined_var_down):
    
    
    mu = param_array[0]
    mu_0 = param_array[1]
    mu_1 = param_array[2]
    mu_2 = param_array[3]
    
    # Here is where I would like to add a pmap decorator to parallelize the loop over the rows of param_val 
    def calculate_variations(carry, param_val):
        
        # All three of this variables have (n_devices, *) shape - I would like to parallelize accross n_devices so each iteration runs over 1 CPU core in parallel
        param, combined_var_up_elem, combined_var_down_elem = param_val
        
        combined_var_array_temp = carry
    
        combined_var_array_temp *= jax.lax.cond(jnp.abs(param)<=1.0, 
                                                poly_interp, 
                                                exp_interp, 
                                                (param, combined_var_up_elem,\
                                                 combined_var_down_elem, tuple_init_alt))
        
        return combined_var_array_temp, None
    
    # Prepare xs for scan
    param_vals = param_array[4:]
    combined_var_up_elems = combined_var_up[:len(param_array)-4]
    combined_var_down_elems = combined_var_down[:len(param_array)-4]
    xs = (param_vals, combined_var_up_elems, combined_var_down_elems)
 
    # The lax.scan loop that I want to use jax.pmap on
    combined_var_array, _ = jax.lax.scan(calculate_variations, tuple_init_alt, xs)
    
    ratio_x = jax.lax.slice(combined_var_array, (0,),(len_x,))
    ratio_y = jax.lax.slice(combined_var_array, (len_x,),(len_x+len_y,))
    
    return sigma(mu, mu_0, mu_1, mu_2, ratio_x, ratio_y)

param_array = np.ones((n_devices+4,))
print(func_optimized_scan(param_array, combined_var_up, combined_var_down))

[JaySandesara]

Hello,

I think I was needlessly confusing myself. With pmap, I dont need to use lax.scan. The problem is resolved simply by pmapping the interpolation functions:

import os, sys, importlib, glob

import numpy as np
import pandas as pd
import math
pd.options.mode.chained_assignment = None
import matplotlib.pyplot as plt
import multiprocessing

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count={}".format(
    multiprocessing.cpu_count()
)

import jax
from jax import numpy as jnp
from functools import partial
from jax import pmap

jax.config.update('jax_platform_name', 'cpu')

platform = jax.lib.xla_bridge.get_backend().platform.casefold()
print("Platform: ", platform)

n_devices = jax.local_device_count()

#Random inputs

arr_ratio_up_x = np.ones((n_devices, 341749))
arr_ratio_down_x = np.ones((n_devices, 341749))

arr_ratio_up_y = np.ones((n_devices, 34179))
arr_ratio_down_y = np.ones((n_devices, 34179))

ones_x = np.ones((arr_ratio_up_x.shape[1],))
ones_y = np.ones((arr_ratio_up_y.shape[1],))

tuple_init = (ones_x,ones_y)

combined_var_up = jnp.column_stack((arr_ratio_up_x, arr_ratio_up_y))
combined_var_down = jnp.column_stack((arr_ratio_down_x, arr_ratio_down_y))

len_x = arr_ratio_up_x.shape[1]
len_y = arr_ratio_up_y.shape[1]

tuple_init_alt = np.ones_like(combined_var_up)

#Interpolation functions - Ignore for this example

@jax.pmap
def poly_interp(tuple_input):

    alpha, varUp, varDown, varNom = tuple_input

    pow_up = jnp.divide(varUp,varNom)
    pow_down = jnp.divide(varDown,varNom)
    logHi = jnp.log(jnp.clip(pow_up,1e-7,None))
    logLo = jnp.log(jnp.clip(pow_down,1e-7,None))
    logHi = jnp.log(pow_up)
    logLo = jnp.log(pow_down)
    pow_up_log = jnp.multiply(pow_up,logHi)
    pow_down_log = -jnp.multiply(pow_down,logLo)
    pow_up_log2 = jnp.multiply(pow_up_log,logHi)
    pow_down_log2 = -jnp.multiply(pow_down_log,logLo)

    S0 = (pow_up+pow_down)/2.0
    A0 = (pow_up-pow_down)/2.0
    S1 = (pow_up_log+pow_down_log)/2.0
    A1 = (pow_up_log-pow_down_log)/2.0
    S2 = (pow_up_log2+pow_down_log2)/2.0
    A2 = (pow_up_log2-pow_down_log2)/2.0

    a1 = (15*A0-7*S1+A2)/8.0
    a2 = (-24+24*S0-9*A1+S2)/8.0
    a3 = (-5*A0+5*S1-A2)/4.0
    a4 = (12-12*S0+7*A1-S2)/4.0
    a5 = (3*A0-3*S1+A2)/8.0
    a6 = (-8+8*S0-5*A1+S2)/8.0

    return 1.0 + alpha * (a1 + alpha * ( a2 + alpha * ( a3 + alpha * ( a4 + alpha * ( a5 + alpha * a6 ) ) ) ) )

@jax.pmap
def exp_interp(tuple_input):

    alpha, varUp, varDown, varNom = tuple_input

    return jnp.where(alpha>1.0, (varUp/varNom)**alpha, (varDown/varNom)**(-alpha))


poly_interp((jnp.ones(n_devices,).reshape(n_devices,1), combined_var_up, combined_var_down,tuple_init_alt))