Torch to tensorflow
Is there a way to convert pytorch code to TensorFlow? While I am a little familiar with TensorFlow, totally new to pytorch. For example,
def get_variation_uncertainty(prediction_score_vectors: List[torch.tensor], matrix_size: Tuple) -&gt; Tuple[np.ndarray, np.ndarray, np.ndarray]:

    prediction_score_vectors = torch.stack(tuple(prediction_score_vectors))

    wt_var = np.var(np.sum(prediction_score_vectors[:, :, 1:].cpu().numpy(), axis=2), axis=0).reshape(matrix_size) * 100
    tc_var = np.var(np.sum(prediction_score_vectors[:, :, [1, 3]].cpu().numpy(), axis=2), axis=0).reshape( matrix_size) *100
    et_var = np.var(prediction_score_vectors[:, :, 3].cpu().numpy(), axis=0).reshape(matrix_size) * 100

    return wt_var.astype(np.uint8), tc_var.astype(np.uint8), et_var.astype(np.uint8)

how can I get the TensorFlow equivalent to the above code?

Per the comment, I would recommend using more tf functions to improve the performance and reduce the amount of GPU-CPU communication necessary. Here is an example
@tf.function
def get_variation_uncertainty_tf(prediction_score_vectors, matrix_size):
    prediction_score_vectors = tf.stack(prediction_score_vectors)    
    wt_var_tmp = tf.math.square(tf.math.reduce_std(tf.reduce_sum(prediction_score_vectors[:, :, 1:], axis=2), axis=0))
    # Two steps because that was getting long
    wt_var = tf.reshape(wt_var_tmp, matrix_size) * 100

    tc_var_tmp = tf.math.square(tf.math.reduce_std(prediction_score_vectors[:, :, 1] + prediction_score_vectors[:, :, 3], axis=0))
    tc_var = tf.reshape(tc_var_tmp, matrix_size) * 100

    et_var_tmp = tf.math.square(tf.math.reduce_std(prediction_score_vectors[:, :, 3], axis=0))
    et_var = tf.reshape(et_var_tmp, matrix_size) * 100
    return tf.cast(wt_var, dtype=tf.uint8), tf.cast(tc_var, dtype=tf.uint8), tf.cast(et_var, dtype=tf.uint8)
    # if you need to return np arrays, do that instead of casting, i.e. (wt_var.numpy()).astype(np.uint8)


Tested and works here, although choosing which method will depend heavily on the shape of your data, feel free to try changing the shape to estimate which is best. In my testing the mostly numpy code is actually better unless you have huge dimensions or will be running it in batches.
https://colab.research.google.com/drive/1miOG6FV9MInanwwQxkYeSXVYirVeUh1r?usp=sharing
