This is a minimal example of the real larger problem I am facing. Consider the function below:
import jax.numpy as jnp
def test(x):
return jnp.sum(x)
I tried to vectorize it by:
v_test = jax.vmap(test)
My inputs to test look like:
x1 = jnp.array([1,2,3])
x2 = jnp.array([4,5,6,7])
x3 = jnp.array([8,9])
x4 = jnp.array([10])
and my input to v_test is:
x = [x1, x2, x3, x4]
If I try:
v_test(x)
I get the error below:
ValueError: vmap got inconsistent sizes for array axes to be mapped:
the tree of axis sizes is:
([3, 4, 2, 1],)
Is there a way to vectorize test over a list of unequal length arrays?
I could avoid this by padding so the arrays have the same length, however, padding is not desired.
JAX does not support ragged arrays, (i.e. arrays in which each row has a different number of elements) so there is currently no way to use vmap for this kind of data. Your best bet is probably to use a Python for loop:
y = [test(xi) for xi in x]
Alternatively, you might be able to express the operation you have in mind in terms of segment_sum or similar operations. For example:
segments = jnp.concatenate([i * jnp.ones_like(xi) for i, xi in enumerate(x)])
result = jax.ops.segment_sum(jnp.concatenate(x), segments)
print(result)
# [ 6 22 17 10]
Another possibility is to pad the input arrays so that they can fit into a standard, non-ragged 2D array.
In the model I want to build, there are two placeholders
x = tf.placeholder('float32', shape=[1000, 10])
tags = tf.placeholder('int32', shape=[1000, 1])
(1000 is just the number of examples)
x holds the inputs to neural networks, tags determines which one of the three neural networks will be used to compute the output.
w1 = tf.get_variable('w1', [10, 1], tf.truncated_normal_initializer())
w2 = tf.get_variable('w2', [10, 1], tf.truncated_normal_initializer())
w3 = tf.get_variable('w3', [10, 1], tf.truncated_normal_initializer())
def nn_1(): return tf.matmul(x, w1)
def nn_2(): return tf.matmul(x, w2)
def nn_3(): return tf.matmul(x, w3)
I want to find an elegant way to implement a TensorFlow graph which can compute the output of an x given its tag.
[x1, x2, x3, ..., xn]
[1, 2, 3, ..., 1]
[nn_1(x), nn_2(x), nn_3(x), ..., nn_1(x)]
If x and tags are not arrays, I can implement it with tf.case, for example,
a = tf.placeholder('int32')
b = tf.placeholder('int32')
result = tf.case(
{
tf.equal(b, 1): a + 1,
tf.equal(b, 2): a + 2
})
But I have no idea how to do when x and tags are arrays.
You can use some math trick to do the job.
let's say you want to implement the code you implemented but with a and b being arrays.
First, you compute an array of condition.
This would be the condition that must be true in order to apply the operation.
Typically conditions use "less", "equal", "greater" operation or a logical composition of those.
You can use tf.bitwise or tf.math.logical* for logical operation and tf.math for the others.
The condition must be a boolean array. 1 if the condition is true, 0 if false.
After that, you initialize the result array with the default value (what is in the "else" statement)
To apply the condition you simply multiply the condition array with the value you want to assign.
The code would be something like this.
//default value
result = tf.zeros(tf.shape(a)[0])
condition = tf.equal(b, index)
condition = tf.cast(condition, tf.float32)
result = tf.multiply(condition, a) + index
If you want to use tag as index of the functions array you need to use a 2d array. Create a 2d array of all possible combination nn X x.
This array will contain nn_j(x[i]) for each i,j couple.
To do this you need to create an array x X nn X 2 array.
First, expand x and create an array with x X nn array
if your x is x=[0,2,1] and len(n) = 2 then you need to have x_nn = [[0,0], [2,2], [1,1]].
nn_x = x
nn_x = tf.expand_dims(nn_x,0)
nn_x = tf.tile(nn_x, [len(nn), 1])
Then you create a 2d array with the same shape having the index of nn.
For the arrey used early index2d = [[0,1],[0,1],[0,1]]
index = tf.linspace(0,len(nn)-1)
index2d = tf.expand_dims(index,0)
index2d = tf.tile(index2d, tf.shape(x)[0])
Then you need to stack these and arrays, move the first dimension at the last place, then flat along axis 0 and 1.
In this way you will have map2d = [[0,0],[0,1],[2,0],[2,1],[1,0],[1,1]]
For each couple the first is the value of x, second is the index of the nn
Then you map this 2d array using the tf.map_fn function. write something like
tf.map_fn(t => [nn[t[1]](t[0]), t[1]], map2d)
Now you have all possible value of nn for each x
At this point, you can reshape back map2d compare map2d[:,:,1] with you tag and select the one that is equal.
#reshape map2d
# ...
# transform tag
tag2d = tag
tag2d = tf.expand_dims(tag2d,0)
tag2d = tf.tile(tag2d, [len(nn), 1])
result = tf.equal(tag, map2d[:,:,1])
the result will have only one non-zero value for each column
result = tf.multiply(result, map2d[:,:,0])
result = tf.reduce_max(result, [1])
I didn't try the code, but the mechanism should work.
Hope this help
I have some 3D and 4D arrays that I want to mutate. Indexing the arrays looks like:
3D: array[time][x][y]
4D: array[time][z][x][y]
I want to do multiple types of data mutation and have been writing a lot of for-in loops. Examples of data mutation:
Replace 0. values with NaN
Create cumulative array (time = 0, x = x0, y = y0 -> time = 1, x = x0 + x1, y = y0 + y1) etc.
Use 2 arrays and combine them (sqrt(xa * xa + xb * xb)))
Numpy can be used for example to do (1) using:
for i in range(len(data)):
if len(data[i].shape) == 3:
for z in range(len(data[i]):
data[i][z][data[i][z] == 0.] = np.NaN
else:
data[i][data[i] == 0.] = np.NaN
But I have to write for loops every time I encounter such a problem and have to write the data mutation 2 times. Once for the 4D array and once for the 3D array. For-loops is a feature to overcome, not having to write if else statements and writing the mutation twice is almost a must. If I somehow or someone else changes the first part but forgets to change the second part the code becomes bugged.
Is there a way to for example enumerate the [z][x][y] but if the array is only [x][y] just do it once as if there was a single z-index like [0][x][y]?
I am having the following numpy arrays:
import numpy as np
y2 = np.array([[0.2,0.1,0.8,0.4],[0.4,0.2,0.5,0.1],[0.4,0.2,0.5,0.1]])
y1 = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1]])
What I am trying to do is to get the position of y1 compared to y2. To be more clear: y1 is the label data and y2 is the predicted data and I want to see in which rank position an algorithm predicted compared with the real data.
I am doing the following:
counter = 0
indexes2 = []
indexes = np.where(y1)[1]
sorted_values = np.argsort(-y2)
for value in sorted_values:
indexes2.append(np.where(value==indexes[counter])[0][0] + 1)
counter += 1
b = np.array(indexes2)
The output is correct:
>>> b
>>> array([2, 2, 3], dtype=int64)
But, I am pretty sure that there is a more elegant way of doing and more optimized. Any hint?
Vectorize the nested loop
We could get rid of the loop by making use of broadcasting -
b = (sorted_values == indexes[:,None]).argmax(1)+1
Some Improvement
For performance, we could optimize the computation of indexes, like so -
indexes = y1.argmax(1)
Bigger Improvement
Additionally, we could optimize on sorted_values computation by avoiding the negation of y2, by doing -
sorted_values2 = np.argsort(y2)
Then, compute b by using broadcasted comparsion as done earlier and subtract the argmax indices from the length of each row. This in effect does the descending ordering along each row as done in the posted question, where we had negation of argsort.
Thus, the final step would be -
b = y2.shape[1] - (sorted_values2 == indexes[:,None]).argmax(1)
I think I missed something somewhere. I filled a numpy array using two for loops (x and y) and a function based on the x,y position. The only problem is that the value of the array always ends in zero irregardless of the size of the array.
thetamap = numpy.zeros(36, dtype=float)
thetamap.shape = (6, 6)
for y in range(0,5):
for x in range(0,5):
thetamap[x][y] = x+y
print thetamap
range(0, 5) produces 0, 1, 2, 3, 4. The endpoint is always omitted. You want simply range(6).
Better yet, use the awesome power of NumPy to make the array in one line:
thetamap = np.arange(6) + np.arange(6)[:,None]
This makes a row vector and a column vector, then adds them together using NumPy broadcasting to make a matrix.