how can i do with vmap todo onehot_fun.

[virgo777]

sorry for my poor english!

i want to use vmap do my own function one_hot,code is here:

def xiaohua_one_hot(idx):
    num_classes = 2
    idx_one_hot = [0] * num_classes
    idx_one_hot[idx] = 1
    return idx_one_hot

when i use

arr = jax.numpy.array([0,1])

for i in arr:
    print(xiaohua_one_hot(i))

it work.but when i use it with vmap.like:

jax.vmap(xiaohua_one_hot,in_axes=(0,None),out_axes=0)(arr)

it error ,i dont know how to finish it.

sorry for poor english

[jakevdp]

Hi - thanks for the question. There are two problems with your implementation

1. jax.vmap can map functions that return arrays; your function returns a list.*
2. jax.vmap will trace your function, meaning that the inputs cannot be used for indexing operations on Python containers; instead you can do indexing operations on JAX arrays (see How To Think In JAX for some background on what tracing means).

Fixing these two problems, I can rewrite your code this way to make it compatible with vmap:

def one_hot(idx):
  num_classes = 2
  idx_one_hot = jax.numpy.zeros(num_classes)
  return idx_one_hot.at[idx].set(1)

arr = jax.numpy.array([0,1])

for i in arr:
    print(one_hot(i))
# [1. 0.]
# [0. 1.]

print(jax.vmap(one_hot,in_axes=0,out_axes=0)(arr))
# [[1. 0.]
#  [0. 1.]]

See this doc page for more information on the ndarray.at I used in place of your list index assignment.

* vmap will also work with functions that return lists, but will treat them as pytrees an map each return element individually, which is probably not what you intended here.

[virgo777]

thanks a lot for your help