body_fun output and input must have identical types...

[jecampagne]

Hello
I have a nasty crash that I do not manage to tackle, I ask for a kind help.

Here a snippet with

1. a function that compute the ratio of two polynomials

import jax
from jax.config import config
config.update("jax_enable_x64", True)

import jax.scipy as jsc
import jax.numpy as jnp
import matplotlib.pyplot as plt

def _evaluate_rational(z, num, denom):
    """ Compute Num(z)/Denom(z)
    """
    count = len(num)

    def true_fn_update(z):

      def body_true(val1):
        # decode val1
        s1, s2, i = val1 
        s1 *= z; s2 *= z; s1 += num[i]; s2 += denom[i];
        return s1, s2, i-1

      def cond_true(val1):
        s1, s2, i = val1 
        return i>=0

      val1_init = (num[-1], denom[-1], count-2)
      s1, s2, _ = jax.lax.while_loop(cond_true, body_true, val1_init)

      return s1/s2

    def false_fn_update(z):
      def body_false(val1):
        # decode val1
        s1, s2, i = val1 
        s1 *= z; s2 *= z; s1 += num[i]; s2 += denom[i];
        return s1, s2, i+1

      def cond_false(val1):
        s1, s2, i = val1 
        return i<count

      val1_init = (num[0], denom[0],1)
      s1, s2, _ = jax.lax.while_loop(cond_false, body_false, val1_init)

      return s1/s2

    return jnp.where(z<=1, true_fn_update(z), false_fn_update(1/z))  

For instance

def Jtest(x):
  P1 = jnp.array([1.,-2.])
  Q1 = jnp.array([1.,2.])
  return  _evaluate_rational(x, P1, Q1)

def Ratio(x):
  return (1-2*x)/(1+2*x)

xvar=jnp.linspace(0,10.,50)
plt.plot(xvar,jax.vmap(Jtest)(xvar))
plt.plot(xvar,Ratio(xvar),ls="--")

2. a Quadrature Integrator

from functools import partial
from jax import grad, jit, vmap
class ClenshawCurtisQuad:
    """
        Clenshaw-Curtis quadrature computed by FFT
    """

    def __init__(self,order=5):
        # 2n-1 quad
        self._order = jnp.int64(2*order-1)
        self._absc, self._absw, self._errw = self.ComputeAbsWeights()
        self.rescaleAbsWeights()
    
    def __str__(self):
        return f"xi={self._absc}\n wi={self._absw}\n errwi={self._errw}"
    
    @property
    def absc(self):
        return self._absc
    @property
    def absw(self):
        return self._absw
    @property
    def errw(self):
        return self._errw
        
    def ComputeAbsWeights(self):
        x,wx = self.absweights(self._order)
        nsub = (self._order+1)//2
        xSub, wSub = self.absweights(nsub)
        errw = jnp.array(wx, copy=True)                           # np.copy(wx)
        errw=errw.at[::2].add(-wSub)             # errw[::2] -= wSub
        return x,wx,errw
  
    
    def absweights(self,n):

        points = -jnp.cos((jnp.pi * jnp.arange(n)) / (n - 1))

        if n == 2:
            weights = jnp.array([1.0, 1.0])
            return points, weights
            
        n -= 1
        N = jnp.arange(1, n, 2)
        length = len(N)
        m = n - length
        v0 = jnp.concatenate([2.0 / N / (N - 2), jnp.array([1.0 / N[-1]]), jnp.zeros(m)])
        v2 = -v0[:-1] - v0[:0:-1]
        g0 = -jnp.ones(n)
        g0 = g0.at[length].add(n)     # g0[length] += n
        g0 = g0.at[m].add(n)          # g0[m] += n
        g = g0 / (n ** 2 - 1 + (n % 2))

        w = jnp.fft.ihfft(v2 + g)
        ###assert max(w.imag) < 1.0e-15
        w = w.real

        if n % 2 == 1:
            weights = jnp.concatenate([w, w[::-1]])
        else:
            weights = jnp.concatenate([w, w[len(w) - 2 :: -1]])
            
        #return
        return points, weights
    
    def rescaleAbsWeights(self, xInmin=-1.0, xInmax=1.0, xOutmin=0.0, xOutmax=1.0):
        """
            Translate nodes,weights for [xInmin,xInmax] integral to [xOutmin,xOutmax] 
        """
        deltaXIn = xInmax-xInmin
        deltaXOut= xOutmax-xOutmin
        scale = deltaXOut/deltaXIn
        self._absw *= scale
        tmp = jnp.array([((xi-xInmin)*xOutmax
                         -(xi-xInmax)*xOutmin)/deltaXIn for xi in self._absc])
        self._absc=tmp
        
@partial(jit, static_argnums=(0,3))
def quadIntegral(f,a,b,quad):
    a = jnp.atleast_1d(a)
    b = jnp.atleast_1d(b)
    d = b-a
    xi = a[jnp.newaxis,:]+ jnp.einsum('i...,k...->ik...',quad.absc,d)
    fi = f(xi)
    S = d * jnp.einsum('i...,i...',quad.absw,fi)
    return S.squeeze()

with a test

quadIntegral(Ratio,0.,10., ClenshawCurtisQuad(10))  # DeviceArray(-6.95548393, dtype=float64)

which is ok according to Mathematica

Now, the following gives a crash

quadIntegral(jax.vmap(Jtest),0.,10., ClenshawCurtisQuad(10))

UnfilteredStackTrace                      Traceback (most recent call last)

[<ipython-input-62-284cfdbdf9b4>](https://localhost:8080/#) in <module>
----> 1 quadIntegral(jax.vmap(Jtest),0.,10., ClenshawCurtisQuad(10))

14 frames

UnfilteredStackTrace: TypeError: body_fun output and input must have identical types, got
('DIFFERENT ShapedArray(float64[1]) vs. ShapedArray(float64[])', 'DIFFERENT ShapedArray(float64[1]) vs. ShapedArray(float64[])', 'ShapedArray(int64[], weak_type=True)').

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------


The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)

[<ipython-input-44-0eca65f4076f>](https://localhost:8080/#) in true_fn_update(z)
     16 
     17       val1_init = (num[-1], denom[-1], count-2)
---> 18       s1, s2, _ = jax.lax.while_loop(cond_true, body_true, val1_init)
     19 
     20       return s1/s2

TypeError: body_fun output and input must have identical types, got
('DIFFERENT ShapedArray(float64[1]) vs. ShapedArray(float64[])', 'DIFFERENT ShapedArray(float64[1]) vs. ShapedArray(float64[])', 'ShapedArray(int64[], weak_type=True)').

But, if I use a simpler integrator (although less accurate in oscillatory functions) then

def simps(f, a, b, N=128):
    """Approximate the integral of f(x) from a to b by Simpson's rule.
    """
    if N % 2 == 1:
        raise ValueError("N must be an even integer.")
    dx = (b - a) / N
    x = jnp.linspace(a, b, N + 1)
    y = f(x)
    S = dx / 3 * jnp.sum(y[0:-1:2] + 4 * y[1::2] + y[2::2], axis=0)
    return S

then

simps(Ratio, 0., 10., N=128) # DeviceArray(-6.95545874, dtype=float64)
simps(jax.vmap(Jtest), 0., 10., N=128) # idem

So I imagine an interaction between the quadIntegraland the _evaluate_rational functions... but which one???

[jecampagne]

Sorry I found just after the post...., but I think the TypeError is not easy to tackle.

I've found a solution comparing the simps node structure and the corresponding quadIntegral nodes:

def quadIntegral(f,a,b,quad):
    a = jnp.atleast_1d(a)
    b = jnp.atleast_1d(b)
    d = b-a
    xi = a[jnp.newaxis,:]+ jnp.einsum('i...,k...->ik...',quad.absc,d)
    xi = xi.squeeze()                                      ######### <-------------------------- HERE
    print("CC  : xi = ",xi)
    fi = f(xi)
    S = d * jnp.einsum('i...,i...',quad.absw,fi)
    return S.squeeze()