Package: kernex
Differentiable Stencil computations in JAXInstallation|Description|Quick example|Function mesh|More Examples|Benchmarking🛠️ Installationpipinstallkernex📖 DescriptionKernex extendsjax.vmap/jax.lax.map/jax.pmapwithkmapandjax.lax.scanwithkscanfor general stencil computations.The prime motivation for this package is to blend the solution process of PDEs into a NN setting.⏩ Quick Examplekmapkscanimportkernexaskeximportjax.numpyasjnp@kex.kmap(kernel_size=(3,))defsum_all(x):returnjnp.sum(x)>>>x=jnp.array([1,2,3,4,5])>>>print(sum_all(x))[6912]importkernexaskeximportjax.numpyasjnp@kex.kscan(kernel_size=(3,))defsum_all(x):returnjnp.sum(x)>>>x=jnp.array([1,2,3,4,5])>>>print(sum_all(x))>>>[61322]`jax.vmap` is used to sum each window content.`lax.scan` is used to update the array and the window sum is calculated sequentially.
the first three rows represents the three sequential steps used to get the solution in the last row.🕸️ Function mesh conceptThe objective is to applyf(x) = x^2  at index=0  and f(x) = x^3 at  index=(1,10)To achieve the following operation withjax.lax.switch, we need a list of 10 functions correspoing to each cell of the example array.
For this reason , kernex adopts a modified version ofjax.lax.switchto reduce the number of branches required.# function applies x^2 at boundaries, and applies x^3 to to the interior┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐f=│x^2│x^3│x^3│x^3│x^3│x^3│x^3│x^3│x^3│x^3│└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐f(│1│2│3│4│5│6│7│8│9│10│)=└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐│1│8│27│64│125│216│343│512│729│1000│└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘# Gradient of this function┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐df/dx=│2x│3x^2│3x^2│3x^2│3x^2│3x^2│3x^2│3x^2│3x^2│3x^2│└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐df/dx(│1│2│3│4│5│6│7│8│9│10│)=└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐│2│12│27│48│75│108│147│192│243│300│└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘Function meshArray equivalentF=kex.kmap(kernel_size=(1,))F[0]=lambdax:x[0]**2F[1:]=lambdax:x[0]**3array=jnp.arange(1,11).astype('float32')print(F(array))>>>[1.,8.,27.,64.,125.,...216.,343.,512.,729.,1000.]print(jax.grad(lambdax:jnp.sum(F(x)))(array))>>>[2.,12.,27.,48.,75.,...108.,147.,192.,243.,300.]defF(x):f1=lambdax:x**2f2=lambdax:x**3x=x.at[0].set(f1(x[0]))x=x.at[1:].set(f2(x[1:]))returnxarray=jnp.arange(1,11).astype('float32')print(F(array))>>>[1.,8.,27.,64.,125.,...216.,343.,512.,729.,1000.]print(jax.grad(lambdax:jnp.sum(F(x)))(array))>>>[2.,12.,27.,48.,75.,...108.,147.,192.,243.,300.]Additionally , we can combine the function mesh concept with stencil computation for scientific computing.
See Linear convection inMore examplessection🔢 More examples1️⃣ Convolution operationimportjaximportjax.numpyasjnpimportkernexaskex@jax.jit@kex.kmap(kernel_size=(3,3,3),padding=('valid','same','same'))defkernex_conv2d(x,w):# JAX channel first conv2d with 3x3x3 kernel_sizereturnjnp.sum(x*w)2️⃣ Laplacian operation# see also# https://numba.pydata.org/numba-doc/latest/user/stencil.html#basic-usageimportjaximportjax.numpyasjnpimportkernexaskex@kex.kmap(kernel_size=(3,3),padding='valid',relative=True)# `relative`= True enables relative indexingdeflaplacian(x):return(0*x[1,-1]+1*x[1,0]+0*x[1,1]+1*x[0,-1]+-4*x[0,0]+1*x[0,1]+0*x[-1,-1]+1*x[-1,0]+0*x[-1,1])# apply laplacian>>>print(laplacian(jnp.ones([10,10])))DeviceArray([[0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.]],dtype=float32)3️⃣ Get Patches of an arrayimportjaximportjax.numpyasjnpimportkernexaskex@kex.kmap(kernel_size=(3,3),relative=True)defidentity(x):# similar to numba.stencil# this function returns the top left cell in the padded/unpadded kernel view# or center cell if `relative`=Truereturnx[0,0]# unlike numba.stencil , vector output is allowed in kernex# this function is similar to# `jax.lax.conv_general_dilated_patches(x,(3,),(1,),padding='same')`@jax.jit@kex.kmap(kernel_size=(3,3),padding='same')defget_3x3_patches(x):# returns 5x5x3x3 arrayreturnxmat=jnp.arange(1,26).reshape(5,5)>>>print(mat)[[12345][678910][1112131415][1617181920][2122232425]]# get the view at array index = (0,0)>>>print(get_3x3_patches(mat)[0,0])[[000][012][067]]4️⃣ Linear convection$\Large {\partial u \over \partial t} + c {\partial u \over \partial x} = 0$$\Large u_i^{n} = u_i^{n-1} - c \frac{\Delta t}{\Delta x}(u_i^{n-1}-u_{i-1}^{n-1})$Problem setupStencil viewimportjaximportjax.numpyasjnpimportkernexaskeximportmatplotlib.pyplotasplt# see https://nbviewer.org/github/barbagroup/CFDPython/blob/master/lessons/01_Step_1.ipynbtmax,xmax=0.5,2.0nt,nx=151,51dt,dx=tmax/(nt-1),xmax/(nx-1)u=jnp.ones([nt,nx])c=0.5# kscan moves sequentially in row-major order and updates in-place using lax.scan.F=kernex.kscan(kernel_size=(3,3),padding=((1,1),(1,1)),named_axis={0:'n',1:'i'},# n for time axis , i for spatial axis (optional naming)relative=True)# boundary condtion as a functiondefbc(u):return1# initial condtion as a functiondefic1(u):return1defic2(u):return2deflinear_convection(u):return(u['i','n-1']-(c*dt/dx)*(u['i','n-1']-u['i-1','n-1']))F[:,0]=F[:,-1]=bc# assign 1 for left and right boundary for all t# square wave initial conditionF[:,:int((nx-1)/4)+1]=F[:,int((nx-1)/2):]=ic1F[0:1,int((nx-1)/4)+1:int((nx-1)/2)]=ic2# assign linear convection function for# interior spatial location [1:-1]# and start from t>0  [1:]F[1:,1:-1]=linear_convectionkx_solution=F(jnp.array(u))plt.figure(figsize=(20,7))forlineinkx_solution[::20]:plt.plot(jnp.linspace(0,xmax,nx),line)5️⃣ Gaussian blurimportjaximportjax.numpyasjnpimportkernexaskexdefgaussian_blur(image,sigma,kernel_size):x=jnp.linspace(-(kernel_size-1)/2.0,(kernel_size-1)/2.0,kernel_size)w=jnp.exp(-0.5*jnp.square(x)*jax.lax.rsqrt(sigma))w=jnp.outer(w,w)w=w/w.sum()@kex.kmap(kernel_size=(kernel_size,kernel_size),padding="same")defconv(x):returnjnp.sum(x*w)returnconv(image)6️⃣ Depthwise convolutionimportjaximportjax.numpyasjnpimportkernexaskex@jax.jit@jax.vmap@kex.kmap(kernel_size=(3,3),padding=('same','same'))defkernex_depthwise_conv2d(x,w):# Channel-first depthwise convolution # jax.debug.print("x=\n{a}\nw=\n{b} \n\n",a=x, b=w)returnjnp.sum(x\*w)h,w,c=5,5,2k=3x=jnp.arange(1,h*w*c+1).reshape(c,h,w)w=jnp.arange(1,k*k*c+1).reshape(c,k,k)print(kernex_depthwise_conv2d(x,w))</summary>7️⃣ Maxpooling2D and Averagepooling2D@jax.vmap# vectorize over the channel dimension@kex.kmap(kernel_size=(3,3),strides=(2,2))defmaxpool_2d(x):# define the kernel for the Maxpool operation over the spatial dimensionsreturnjnp.max(x)@jax.vmap# vectorize over the channel dimension@kex.kmap(kernel_size=(3,3),strides=(2,2))defavgpool_2d(x):# define the kernel for the Average pool operation over the spatial dimensionsreturnjnp.mean(x)8️⃣ Runge-Kutta integration# lets solve dydt = y, where y0 = 1 and y(t)=e^t# using Runge-Kutta 4th order method# f(t,y) = yimportjax.numpyasjnpimportmatplotlib.pyplotaspltimportkernexaskext=jnp.linspace(0,1,5)y=jnp.zeros(5)x=jnp.stack([y,t],axis=0)dt=t[1]-t[0]# 0.1f=lambdatn,yn:yndefic(x):""" initial condition y0 = 1 """return1.defrk4(x):""" runge kutta 4th order integration step """# ┌────┬────┬────┐      ┌──────┬──────┬──────┐# │ y0 │*y1*│ y2 │      │[0,-1]│[0, 0]│[0, 1]│# ├────┼────┼────┤ ==>  ├──────┼──────┼──────┤# │ t0 │ t1 │ t2 │      │[1,-1]│[1, 0]│[1, 1]│# └────┴────┴────┘      └──────┴──────┴──────┘t0=x[1,-1]y0=x[0,-1]k1=dt*f(t0,y0)k2=dt*f(t0+dt/2,y0+1/2*k1)k3=dt*f(t0+dt/2,y0+1/2*k2)k4=dt*f(t0+dt,y0+k3)yn_1=y0+1/6*(k1+2*k2+2*k3+k4)returnyn_1F=kex.kscan(kernel_size=(2,3),relative=True,padding=((0,1)))# kernel size = 3F[0:1,1:]=rk4F[0,0]=ic# compile the solversolver=jax.jit(F.__call__)y=solver(x)[0,:]plt.plot(t,y,'-o',label='rk4')plt.plot(t,jnp.exp(t),'-o',label='analytical')plt.legend()