TensorFlow | How to use tf.stack() in 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