I need help figuring out how to use jax.dynamic_slice_in_dim

[richinex]

Hello I am trying to rewrite my function using jax.lax.dynamic_slice_in_dim but I can not get it to work. I have shown the two examples below and would like to know what I am doing wrong.
The problem occurs when I set an element in smf to inf. Then I get different values.

Example 1 (What I have always used):

import numpy as np

num_rows = 5
num_cols = 20
smf = np.array([np.inf, 0.1, 0.1, 0.1, 0.1])
par_init = np.array([1.0,2.0,3.0,4.0,5.0])
lb = np.array([0.1, 0.1, 0.1, 0.1, 0.1])
ub = np.array([10.0, 10.0, 10.0, 10.0, 10.0])
par = np.broadcast_to(par_init[:,None],(num_rows,num_cols))

kvals = np.where(np.isinf(smf), 1, num_cols)
kvals = np.insert(kvals, 0, 0)
kvals = np.cumsum(kvals)

par0_col = np.zeros(num_rows*num_cols - (num_cols-1) * np.sum(np.isinf(smf)))
lb_col = np.zeros(num_rows*num_cols - (num_cols-1) * np.sum(np.isinf(smf)))
ub_col = np.zeros(num_rows*num_cols- (num_cols-1) * np.sum(np.isinf(smf)))



for i in range(num_rows):
    par0_col[kvals[i]:kvals[i+1]] = par[i, :kvals[i+1]-kvals[i]]
    lb_col[kvals[i]:kvals[i+1]] = lb[i]
    ub_col[kvals[i]:kvals[i+1]] = ub[i]



par_log = np.log10((par0_col - lb_col) / (1 - par0_col / ub_col))

arr_1 = np.zeros(shape = (num_rows, num_cols))
arr_2 = np.zeros(shape = (num_rows, num_cols))

for i in range(num_rows):


    arr_1[i, :] = (par_log[kvals[i]:kvals[i+1]])
    arr_2[i, :] = (lb_col[kvals[i]:kvals[i+1]] + 10**par_log[kvals[i]:kvals[i+1]]) / (1 + 10**par_log[kvals[i]:kvals[i+1]] / ub_col[kvals[i]:kvals[i+1]])

print(f"arr_1 is \n{arr_1}")

print(f"arr_2 is \n{arr_2}")

# What I expect
# arr_1 is
# [[0.         0.         0.         0.         0.         0.
#   0.         0.         0.         0.         0.         0.
#   0.         0.         0.         0.         0.         0.
#   0.         0.        ]
#  [0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361]
#  [0.61729996 0.61729996 0.61729996 0.61729996 0.61729996 0.61729996
#   0.61729996 0.61729996 0.61729996 0.61729996 0.61729996 0.61729996
#   0.61729996 0.61729996 0.61729996 0.61729996 0.61729996 0.61729996
#   0.61729996 0.61729996]
#  [0.81291336 0.81291336 0.81291336 0.81291336 0.81291336 0.81291336
#   0.81291336 0.81291336 0.81291336 0.81291336 0.81291336 0.81291336
#   0.81291336 0.81291336 0.81291336 0.81291336 0.81291336 0.81291336
#   0.81291336 0.81291336]
#  [0.99122608 0.99122608 0.99122608 0.99122608 0.99122608 0.99122608
#   0.99122608 0.99122608 0.99122608 0.99122608 0.99122608 0.99122608
#   0.99122608 0.99122608 0.99122608 0.99122608 0.99122608 0.99122608
#   0.99122608 0.99122608]]
# arr_2 is
# [[1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
#  [2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]
#  [3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3.]
#  [4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4.]
#  [5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5.]]

Example 2 (When I use jax.lax.dynamic_slice_in_dim):

import jax
import jax.numpy as jnp
import numpy as np
jax.config.update("jax_enable_x64", True)

num_rows = 5
num_cols = 20
smf = jnp.array([jnp.inf, 0.1, 0.1, 0.1, 0.1])
par_init = jnp.array([1.0,2.0,3.0,4.0,5.0])
lb = jnp.array([0.1, 0.1, 0.1, 0.1, 0.1])
ub = jnp.array([10.0, 10.0, 10.0, 10.0, 10.0])
par = jnp.broadcast_to(par_init[:,None],(num_rows,num_cols))

kvals = jnp.where(jnp.isinf(smf), 1, num_cols)
kvals = jnp.insert(kvals, 0, 0)
kvals = jnp.cumsum(kvals)

par0_col = jnp.zeros(num_rows*num_cols - (num_cols-1) * jnp.sum(jnp.isinf(smf)))
lb_col = jnp.zeros(num_rows*num_cols - (num_cols-1) * jnp.sum(jnp.isinf(smf)))
ub_col = jnp.zeros(num_rows*num_cols- (num_cols-1) * jnp.sum(jnp.isinf(smf)))



for i in range(num_rows):
    par0_col = par0_col.at[kvals[i]:kvals[i+1]].set(par[i, :kvals[i+1]-kvals[i]])
    lb_col = lb_col.at[kvals[i]:kvals[i+1]].set(lb[i])
    ub_col = ub_col.at[kvals[i]:kvals[i+1]].set(ub[i])



# par_log = jnp.asarray(np.random.lognormal(mean=0.0, sigma=1.0, size=100))

par_log = jnp.log10((par0_col - lb_col) / (1 - par0_col / ub_col))


def get_arr_1(i):
    return (jax.lax.dynamic_slice_in_dim(par_log, kvals[i], num_cols))


def get_arr_2(i):
    return ((jax.lax.dynamic_slice_in_dim(lb_col, kvals[i],num_cols) + jax.lax.dynamic_slice_in_dim(10**par_log, kvals[i],num_cols))/ \
        (1 + jax.lax.dynamic_slice_in_dim(10**par_log, kvals[i],num_cols) / jax.lax.dynamic_slice_in_dim(ub_col, kvals[i],num_cols)))

arr_1 = (jax.vmap(get_arr_1, in_axes=0) (jnp.arange(num_rows)))
arr_2 =(jax.vmap(get_arr_2, in_axes=0) (jnp.arange(num_rows)))


print(f"arr_1 is \n{arr_1}")

print(f"arr_2 is \n{arr_2}")

# What I get when I use jax.lax.dynamic_slice
# arr_1 is
# [[0.         0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361]
#  [0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361 0.37566361 0.37566361 0.37566361 0.37566361
#   0.37566361 0.37566361]
#  [0.61729996 0.61729996 0.61729996 0.61729996 0.61729996 0.61729996
#   0.61729996 0.61729996 0.61729996 0.61729996 0.61729996 0.61729996
#   0.61729996 0.61729996 0.61729996 0.61729996 0.61729996 0.61729996
#   0.61729996 0.61729996]
#  [0.81291336 0.81291336 0.81291336 0.81291336 0.81291336 0.81291336
#   0.81291336 0.81291336 0.81291336 0.81291336 0.81291336 0.81291336
#   0.81291336 0.81291336 0.81291336 0.81291336 0.81291336 0.81291336
#   0.81291336 0.81291336]
#  [0.99122608 0.99122608 0.99122608 0.99122608 0.99122608 0.99122608
#   0.99122608 0.99122608 0.99122608 0.99122608 0.99122608 0.99122608
#   0.99122608 0.99122608 0.99122608 0.99122608 0.99122608 0.99122608
#   0.99122608 0.99122608]]
# arr_2 is
# [[1. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]
#  [2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]
#  [3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3.]
#  [4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4.]
#  [5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5.]]