TensorFlow - Creating a Tensor out of a stencil

Question

While searching for this problem I haven´t found something related, so please let me know if something similar already exists.

So here is the problem. Imagine you have a stencil Tensor A containing the non-zero values of a column of B. Now I want to generate a Tensor B containing the values of A at the correct position. A should be of dim (n,3) and B of dim (2n+1,n). The different dimension arise from different Finite Element Grids. For example let n=3 say we have A

A = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

and want to create a Tensor B out of A like

B[0,0]=A[0,0], B[1,0]=A[0,1], B[2,0]=A[0,2], B[2,1]=A[1,0], B[3,1]=A[1,1], B[4,1]=A[1,2], B[4,2]=A[2,0], B[5,2]=A[2,1] and B[6,2]=A[2,2].

So for the example B would be

B = [[1, 0, 0], [2, 0, 0], [3, 4, 0], [0, 5, 0], [0, 6, 7], [0, 0, 8], [0, 0, 9]].

Additionaly this needs to be differentiable, if that is even possible. Since I want to train the Tensor A to some parameters, but use B to calculate a Matrix-vector product for the loss funktion. I won´t use batch learning at this point.

I tried to use the tf.assign() function to assign every single Value for the Tensor B. I am looking for a better option since this is getting really messy for large tensors and you have to run the tf.assign() op before the value is even set.

In numpy I would simply do something like

import numpy as np
A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = np.zeros((7,3))
for i in range(3):
    B[2*i:2*i+3,i] = A[i,:]

is there some similar or even easier way to do this in Tensorflow at which the operation is still differentiable?

It would help if you can clarify what the pattern of B is, and why B has 7 in 0-dim. — zihaozhihao

jdehesa jdehesa · Accepted Answer · 2019-10-21T11:03:08

You can do that with tf.scatter_nd like this (function works for TF 1.x and 2.x):

import tensorflow as tf

def make_tensor(a):
    a = tf.convert_to_tensor(a)
    s = tf.shape(a)
    n = s[0]
    m = s[1]
    out_shape = [1 + (m - 1) * n, n]
    r = tf.expand_dims(tf.range(n), 1)
    idx_row = r * (m - 1) + tf.range(m)
    idx_col = tf.tile(r, [1, m])
    return tf.scatter_nd(tf.stack([idx_row, idx_col], axis=-1), a, out_shape)

print(make_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 4 0]
#  [0 5 0]
#  [0 6 7]
#  [0 0 8]
#  [0 0 9]]

EDIT: If you want, you can modify the function above a bit to choose how much overlap there should be between the columns:

import tensorflow as tf

def make_tensor(a, overlap=1):
    s = tf.shape(a)
    n = s[0]
    m = s[1]
    m2 = m - overlap
    out_shape = [overlap + m2 * n, n]
    r = tf.expand_dims(tf.range(n), 1)
    idx_row = r * m2 + tf.range(m)
    idx_col = tf.tile(r, [1, m])
    return tf.scatter_nd(tf.stack([idx_row, idx_col], axis=-1), a, out_shape)

a = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(make_tensor(a).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 4 0]
#  [0 5 0]
#  [0 6 7]
#  [0 0 8]
#  [0 0 9]]
print(make_tensor(a, overlap=2).numpy())
# [[1 0 0]
#  [2 4 0]
#  [3 5 7]
#  [0 6 8]
#  [0 0 9]]
print(make_tensor(a, overlap=3).numpy())
# [[1 4 7]
#  [2 5 8]
#  [3 6 9]]

Which incidentally allows you to give no overlap or even negative overlap:

print(make_tensor(a, overlap=0).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 0 0]
#  [0 4 0]
#  [0 5 0]
#  [0 6 0]
#  [0 0 7]
#  [0 0 8]
#  [0 0 9]]
print(make_tensor(a, overlap=-1).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 0 0]
#  [0 0 0]
#  [0 4 0]
#  [0 5 0]
#  [0 6 0]
#  [0 0 0]
#  [0 0 7]
#  [0 0 8]
#  [0 0 9]]

TensorFlow - Creating a Tensor out of a stencil

1 Answers