title: Batch size has been changed after multiplying input with a matrix
body: <p>I am writing a custom layer in TensorFlow to calculate PCA. When I feed input with shape(None, 150, 300) to this layer, I receive output with shape (150, 150, 100) with n_components=100.</p>
  <p>How can I keep the batch size = None? The output I want is: (None, 150, 100).</p>
  <pre><code>def call(self, input_data):
              covariance_matrix = tfp.stats.covariance(input_data)
          
              self.eigen_values, self.eigen_vectors = tf.linalg.eig(covariance_matrix)
              self.eigen_values = tf.cast(self.eigen_values, dtype=tf.float32)
              self.eigen_vectors = tf.cast(self.eigen_vectors, dtype=tf.float32)
      
              self.projection_matrix = tf.transpose(tf.transpose(self.eigen_vectors)[:][-self.n_components:])
          
              output_data = tf.matmul(input_data, self.projection_matrix)
              
              return output_data
  </code></pre>