code: def dct(matrix):
      """Type-II discrete cosine transform (DCT).
  
      Args:
        matrix: A tensor of shape [batch, height, width, channel].
  
      Returns:
        A tensor of the same shape as input where DCT is applied on the height and
        width dimensions.
      """
      # matrix: [B, H, W, C]  (batch, height, width, channel)
      input_dim = matrix.shape[1]
      dct_mat = np.zeros((input_dim, input_dim), dtype=np.float32)
      # Constructs the DCT transform matrix.
      for i in range(input_dim):
          for j in range(input_dim):
              dct_mat[i, j] = (
                  2.0 / math.sqrt(2.0 * input_dim)
              ) * math.cos(
                  i
                  * (2.0 * j + 1.0)
                  * math.pi
                  * 1.0
                  / (2.0 * input_dim)
              )
      dct_mat[0, :] = dct_mat[0, :] * 1.0 / math.sqrt(2)
  
      # matrix_t: [B, C, W, H]
      matrix_t = tf.transpose(matrix, perm=[0, 3, 2, 1])
      # Computes the transform. Multiplies dct_mat once for the rows, and once for
      # the columns.
      transformed_1d = tf.matmul(matrix_t, dct_mat.T)
      transformed_2d = tf.matmul(dct_mat, transformed_1d)
  
      # result: [B, H, W, C]
      return tf.transpose(transformed_2d, perm=[0, 3, 2, 1])
  
file_path: neural-structured-learning/research/multi_representation_adversary/multi_representation_adversary/transforms.py
imports: import math
  import gin
  import numpy
  import tensorflow.compat.v2
label_desc: Use image node_id 1 for calling a global function with example usage: dct(matrix) and returns: tf