Jit for Array functions

[cpcpython]

Hi List,
I am trying a Tridiagonalization routine in Jax and Python 3.10 (with TPU backend)
(for Tridiagonalization numpy See: https://stackoverflow.com/questions/76852295 )
The working code is attached below.

In order to speed up this I want to add @jit before the function tri() - the motivation is:
"None of the x.at expressions modify the original x; instead they return a modified copy of x. However, inside a jit() compiled function, expressions like x = x.at[idx].set(y) are guaranteed to be applied in-place." as it is shown in:[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html]

In other words, since I used:

v=v.at[i+1].set((H[i+1][i]-alpha)/(2*r))
v=v.at[j].set(H[j][i]/(2*r))

and these 'v' is actually Copies of 'v' if I didnt use @jit as per the above statements (which make them in efficient program, If, say, d > 100).

However when I uncomment @jit it made error (as shown below)
So is any way to JIT this function ?
Any helps will be greatly appreciated

Thanks in advance

Krishna

Error Log ...

[[ 122. 20. -69. 110. 19.]
[ 20. -16. -126. 47. -113.]
[ -69. -126. -82. 187. -53.]
[ 110. 47. 187. -144. 30.]
[ 19. -113. -53. 30. -172.]]
Traceback (most recent call last):
File "/home/gpkmohan_mbcet/Tridiagonalz/j.py", line 56, in
H2=tri(H0)
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 338, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr, attrs_tracked = _python_pjit_helper(
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 178, in _python_pjit_helper
p, args_flat = _infer_params(fun, jit_info, args, kwargs)
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 744, in _infer_params
p, args_flat = _infer_params_impl(
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 641, in _infer_params_impl
jaxpr, consts, out_avals, attrs_tracked = _create_pjit_jaxpr(
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 352, in memoized_fun
ans = call(fun, *args)
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 1288, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts, attrs_tracked = pe.trace_to_jaxpr_dynamic(
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 333, in wrapper
return func(*args, **kwargs)
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2278, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts, attrs_tracked = trace_to_subjaxpr_dynamic(
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/src/interpreters/partial_eval.py", line 2301, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/linear_util.py", line 193, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/gpkmohan_mbcet/Tridiagonalz/j.py", line 41, in tri
alpha= -sgn(H[i+1][i])*norm(H[i+1:,i])
File "/home/gpkmohan_mbcet/Tridiagonalz/j.py", line 19, in sgn
if x>=0:
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/core.py", line 788, in bool
return self.aval._bool(self)
File "/home/gpkmohan_mbcet/.local/lib/python3.10/site-packages/jax/_src/core.py", line 1554, in error
raise TracerBoolConversionError(arg)
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function tri at /home/gpkmohan_mbcet/Tridiagonalz/j.py:38 for jit. This concrete value was not available in Python because it depends on the value of the argument H.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

---

import matplotlib
import matplotlib.pyplot as plt
from jax import config
config.update("jax_enable_x64", True)
import numpy as np
from jax import jit

import jax.numpy as jnp
from jax import random
import time
from jax.lib import xla_bridge

from jax.numpy.linalg import norm
def sgn(x):
    if x>=0:
        return 1
    else:
        return -1
"Lets first generate a random matrix of dimension d"
d=5
np.random.seed(0)
xr=np.random.randint(-111,111,(d,d))
a = xr+ xr.T
H1=(a+a.T)/2
H0=jnp.array(H1)
print(H0)

plt.imshow(H0) #Needs to be in row,col order
plt.savefig("i.png")

# function begin
from functools import partial

@jit
def tri(H):
    for i in range(d-1):  
        alpha= -sgn(H[i+1][i])*norm(H[i+1:,i])
        r= jnp.sqrt(0.5*(alpha**2- H[i+1][i]*alpha))
        v= jnp.zeros((d,1)) # jnp immutable array ....
        v=v.at[i+1].set((H[i+1][i]-alpha)/(2*r))
        for j in range(i+2,d):
            v=v.at[j].set(H[j][i]/(2*r))
        w=jnp.eye(d)-2*jnp.matmul(v,v.T)
        H= jnp.matmul(w,jnp.matmul(H,w))

    return H

#    result_shape = jax.core.ShapedArray(H.shape, H.dtype)
#    return jax.pure_callback(H, result_shape, H)

t1=time.time()
H2=tri(H0)
print(H2)
t2=time.time()

print("Time in sec. = " , t2-t1)

import matplotlib.pyplot as plt
plt.imshow(H1) #Needs to be in row,col order
plt.savefig("f.png")

[RedHat2003]

try

def sgn(x):
    return jnp.where(x >= 0, 1, -1)

[cpcpython]

Thanks a lot...it works!

I wonder Why it works? Perhaps @jit need an Array based sgn() ?

[RedHat2003]

Issue:
When applying @jit to the tri() function, a TracerBoolConversionError is raised due to the use of Python if statements within the function.

Cause:
JAX's JIT compilation leverages the XLA compiler, which requires a static computation graph. Python control flow constructs like if statements depend on runtime values, making them dynamic and incompatible with XLA's expectations.
Thats it 😁

[cpcpython]

Thanx for your explaining. Ciao