Store array of tensors in for loop-CodePudding

I have a @tf.function decorated function. Inside the function, I would like to draw from a distribution and compute some value several times (let's call this proceduce f(x)).

How can I do this in Tensorflow 2.0? I can't use numpy arrays as I would like to use the @tf.function decorator.

A numpy implementation would look like:

reps = 4
store = np.zeros((n, reps))
for i in range(reps):
  store[:, i] = f(x) #f(x) is shape (n,)

The goal would then be to compute the row means of store.

This should be easy but I haven't been able to work out how to do it!

CodePudding user response：

Something like this maybe:

import tensorflow as tf

def f():
  return tf.random.normal((10,))

@tf.function
def store_this():
  reps = 4
  n = 10
  store = tf.zeros((n, reps))
  values = [f() for _ in range(reps)]
  indices = tf.stack([tf.tile(tf.range(n), multiples=[reps]), tf.repeat(tf.range(reps), repeats=n)], axis=-1)
  return tf.tensor_scatter_nd_update(store, indices, tf.reshape(values, [-1]))

store_this()

CodePudding user response：

If f takes a one-dimensional tensor as input, this is a shorter alternative:

@tf.function
def f(x):
  return tf.random.normal((10,))
x = tf.constant([1.0, 2.0])
reps = 4

def store_this(fp, x, reps):
    return tf.transpose(tf.map_fn(fp, tf.tile(tf.expand_dims(x, 0),[reps,1])))
store_this(f, x, reps)