Simplest solution to convert a 1D tensor to a 5x5 tensor-CodePudding

What would be the simplest solution to a convert the tensor:

tensor = tf.constant([4.0])

to a 5x5 tensor, where the main diagonal and the antidiagonal have the scalar 4.0, but the rest of the tensor has the value 0.0 as shown below:

tf.Tensor(
[[4. 0. 0. 0. 4.]
 [0. 4. 0. 4. 0.]
 [0. 0. 4. 0. 0.]
 [0. 4. 0. 4. 0.]
 [4. 0. 0. 0. 4.]], shape=(5, 5), dtype=float32)

CodePudding user response：

Simply use tf.eye:

import tensorflow as tf

tensor = tf.where(tf.greater(tf.reverse(tf.eye(5), axis=[1])   tf.eye(5), 0.0), 1.0, 0.0) * tf.constant([4.0])

tf.Tensor(
[[4. 0. 0. 0. 4.]
 [0. 4. 0. 4. 0.]
 [0. 0. 4. 0. 0.]
 [0. 4. 0. 4. 0.]
 [4. 0. 0. 0. 4.]], shape=(5, 5), dtype=float32)

CodePudding user response：

Here's one way you could do this:


def diag_antidiag(shape):
    """Create an diag and anti-diagonal tensor of ones given `shape`

    Examples
    --------
    >>> diag_antidiag(5)
    (<tf.Tensor: shape=(5, 5), dtype=int32, numpy=
     array([[1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0],
            [0, 0, 0, 1, 0],
            [0, 0, 0, 0, 1]], dtype=int32)>,
     <tf.Tensor: shape=(5, 5), dtype=int32, numpy=
     array([[0, 0, 0, 0, 1],
            [0, 0, 0, 1, 0],
            [0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0],
            [1, 0, 0, 0, 0]], dtype=int32)>)
    """
    diag = tf.linalg.tensor_diag([1] * shape)
    anti = diag[:, ::-1]
    return diag, anti


# create the initial diags
shape = 5
diag, anti = diag_antidiag(shape)

# cast as bool tensors
diag = tf.cast(diag, tf.bool)
anti = tf.cast(anti, tf.bool)

# overlay bool masks
overlay = tf.cast(tf.logical_or(diag, anti), tf.float32)
fours = overlay * 4

# <tf.Tensor: shape=(5, 5), dtype=float32, numpy=
# array([[4., 0., 0., 0., 4.],
#        [0., 4., 0., 4., 0.],
#        [0., 0., 4., 0., 0.],
#        [0., 4., 0., 4., 0.],
#        [4., 0., 0., 0., 4.]], dtype=float32)>

This creates two bool tensors along the diags, overlays them, and then fills them with 4s.

CodePudding user response：

A simplified version of AloneTogether's answer

tf.cast(tf.logical_or(tf.reverse(tf.eye(5, dtype=tf.bool), axis=[1]), tf.eye(5, dtype=tf.bool)), dtype=tf.float32) * 4.0