TensorFlow | How to use tf.stack() in tensorflow

gcptutorials.com TensorFlow

This code snippet is using TensorFlow2.0, if you are using earlier versions of TensorFlow than change the code accordingly.

stack() method of TensorFlow creates a stacked tensor whose rank is one higher than each tensors given in "value" parameter.

Stacking can be done along different "axis" dimensions, refer below code for using tf.stack() method in TensorFlow.

Create tensors and check tensors rank


import tensorflow as tf

print(tf.__version__)

#Create tesnors with tf.random.uniform
tensor1 = tf.random.uniform(shape=[2,2], minval=3, maxval=5)
tensor2 = tf.random.uniform(shape=[2,2], minval=3, maxval=5)
tensor3 = tf.random.uniform(shape=[2,2], minval=3, maxval=5)

#Print tensors and rank
print(tensor1)
print(tf.rank(tensor1))
print(tensor2)
print(tf.rank(tensor2))
print(tensor3)
print(tf.rank(tensor3))

'''Output
2.0.0
tf.Tensor(
[[4.015283  4.2504835]
  [3.3398144 3.6719174]], shape=(2, 2), dtype=float32)
tf.Tensor(2, shape=(), dtype=int32)

tf.Tensor(
[[4.373131  3.4104054]
  [4.9401774 4.645204 ]], shape=(2, 2), dtype=float32)
tf.Tensor(2, shape=(), dtype=int32)

tf.Tensor(
[[3.0604298 3.2270956]
  [4.1364336 3.9556732]], shape=(2, 2), dtype=float32)
tf.Tensor(2, shape=(), dtype=int32)
'''

As we can see from the output tensors are of rank 2, after applying stack() on these tensors output would be a stacked tensor with rank 3.


# Applying tf.stack()
stacked_tensor = tf.stack(values=[tensor1, tensor2, tensor3], axis=0)

print(stacked_tensor)
print(tf.rank(stacked_tensor))
'''
Output
tf.Tensor(
[[[3.539135  4.3329363]
[4.866601  3.1803966]]

[[3.5696192 4.18647  ]
[4.420633  4.330311 ]]

[[3.5565314 4.912521 ]
[4.4694653 4.0235806]]], shape=(3, 2, 2), dtype=float32)
tf.Tensor(3, shape=(), dtype=int32)
'''

   


Below is the complete code snippet.


  import tensorflow as tf

  print(tf.__version__)
  
  #Create tesnors with tf.random.uniform
  tensor1 = tf.random.uniform(shape=[2,2], minval=3, maxval=5)
  tensor2 = tf.random.uniform(shape=[2,2], minval=3, maxval=5)
  tensor3 = tf.random.uniform(shape=[2,2], minval=3, maxval=5)
  
  #Print tensors and rank
  print(tensor1)
  print(tf.rank(tensor1))
  print(tensor2)
  print(tf.rank(tensor2))
  print(tensor3)
  print(tf.rank(tensor3))
  
  
  # Applying tf.stack()
  stacked_tensor = tf.stack(values=[tensor1, tensor2, tensor3], axis=0)
  
  print(stacked_tensor)
  print(tf.rank(stacked_tensor))
  '''
  Output
  2.0.0
  tf.Tensor(
  [[3.539135  4.3329363]
    [4.866601  3.1803966]], shape=(2, 2), dtype=float32)
  tf.Tensor(2, shape=(), dtype=int32)
  
  tf.Tensor(
  [[3.5696192 4.18647  ]
    [4.420633  4.330311 ]], shape=(2, 2), dtype=float32)
  tf.Tensor(2, shape=(), dtype=int32)
  
  tf.Tensor(
  [[3.5565314 4.912521 ]
    [4.4694653 4.0235806]], shape=(2, 2), dtype=float32)
  tf.Tensor(2, shape=(), dtype=int32)
  
  tf.Tensor(
  [[[3.539135  4.3329363]
    [4.866601  3.1803966]]
  
    [[3.5696192 4.18647  ]
    [4.420633  4.330311 ]]
  
    [[3.5565314 4.912521 ]
    [4.4694653 4.0235806]]], shape=(3, 2, 2), dtype=float32)
  tf.Tensor(3, shape=(), dtype=int32)
  '''

Category: TensorFlow