Passing indexes in vmap

[IrisTeng]

The function I am trying to vectorize contains an index as positional argument. Therefore, when I use vmap, I pass a np.array of shape (n, 1) as the list of indexes.

import jax
import jax.numpy as jnp
import numpy as np
from functools import partial
from jax import jit
from sklearn.tree import DecisionTreeRegressor
r_xx = np.random.RandomState(921)
x_train = r_xx.uniform(-2, 2, (20, 15))
f = lambda x: x[:,0] - x[:,1] + x[:,2] + 0.5*x[:,3] + 2*x[:,4]
y_train = f(x_train)
model = DecisionTreeRegressor(random_state=0, max_depth=3).fit(x_train, y_train)

feature_tr0 = r_xx.binomial(1, 0.2, 160).reshape((20, 8))
node_indicator = model.decision_path(x_train)
leaf_id_tr = model.apply(x_train)

# @partial(jit, static_argnums=(1,))
def compute_X_transform(x, sample_id):
    tmp = 0.9
    # obtain ids of the nodes `sample_id` goes through, i.e., row `sample_id`
    node_index = node_indicator.indices[jnp.array(node_indicator.indptr)[sample_id]:
                                        jnp.array(node_indicator.indptr)[sample_id + 1]]
    feature_x = jnp.array(feature_tr0)[sample_id, :]
    for node_id in node_index:
        if jnp.array(leaf_id_tr)[sample_id] == node_id:
            continue
        tmp = tmp**2
    feature_x = tmp * feature_x
    return feature_x

def compute_f(x, sample_id):
    feature_new = compute_X_transform(x, sample_id)
    r_b = np.random.RandomState(921)
    wb = r_b.uniform(-2, 2, (8, ))
    return np.reshape(feature_new @ wb, ())

grad_f = jax.grad(compute_f)
n, d = x_train.shape
grad_train = np.zeros((n, d))
# if I do the for loop, it works
for i in range(n):
    grad_train[i, :] = grad_f(x_train[i, :], i)

# but not using vmap
v_res = jax.vmap(grad_f, in_axes=(0, 0), out_axes=0)(x_train, np.arange(n).reshape((-1,1)))

I found out that there should be index error in

node_index = node_indicator.indices[jnp.array(node_indicator.indptr)[sample_id]:
                                        jnp.array(node_indicator.indptr)[sample_id + 1]]

jax._src.errors.TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int32[1])>with<BatchTrace(level=1/0)> with
  val = DeviceArray([[ 0],
             [ 4],
             [ 8],
             [12],
             [16],
             [20],
             [24],
             [28],
             [32],
             [36],
             [40],
             [44],
             [48],
             [52],
             [56],
             [60],
             [64],
             [68],
             [72],
             [76]], dtype=int32)
  batch_dim = 0

and also bool error in

if jnp.array(leaf_id_tr)[sample_id] == node_id

jax._src.errors.ConcretizationTypeError
Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[1])>with<BatchTrace(level=1/0)> with
  val = DeviceArray([[False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False],
             [False]], dtype=bool)
  batch_dim = 0
The problem arose with the `bool` function. 

But I cannot convert sample_id into int. On the other hand, if I add

@partial(jit, static_argnums=(1,))

before the compute_X_transform function, i.e., keep sample_id static, it throws the unhashable error. Is there a way I can make it to work? I would be greatly appreciative!

[jakevdp]

Your code is missing the definition of rfdt – can you update it so we can reproduce the error? Thanks!

[IrisTeng]

Thanks for your quick reply! I have already updated the code. It should run smoothly now.

[jakevdp]

Fundamentally, the issue is that your function constructs arrays that are not statically sized. For example, the length of node_index changes from iteration to iteration, and so this operation is not compatible with JAX transforms like jit and vmap, which require any created arrays to have static shapes. You won't be able to use this function with JAX transforms until you rewrite its logic with that in mind.

[IrisTeng]

Thanks Jake! Is it possible to convert sample_id (abstract tracer) into int in function compute_X_transform?

[jakevdp]

No, not in vmap. If sample_id is being vmapped over, it is not possible to convert it into an int within the mapped function. vmap works by replacing concrete values with abstract values so that it can effectively evaluate every iteration in a single pass. If you have a function that cannot be abstractly evaluated, you have a function that is incompatible with vmap.

[IrisTeng]

Hi Jake, I rewrote the format in another way and resolved the issue. Thanks!