class_docstr: Define how to apply SVD algorithm.
code: def init_training_weights(self, pretrained_weight: tf.Tensor):
      rank = self.rank
      s, u, v = tf.linalg.svd(pretrained_weight)
  
      if len(pretrained_weight.shape) == 2:
          # FC Layer
          s = s[:rank]
          u = u[:, :rank]
          v = v[:, :rank]
      elif len(pretrained_weight.shape) == 4:
          # Conv2D Layer
          s = s[:, :, :rank]
          u = u[:, :, :, :rank]
          v = v[:, :, :, :rank]
      else:
          raise NotImplementedError(
              "Only for dimension=2 or 4 is supported."
          )
  
      sv = tf.matmul(tf.linalg.diag(s), v, adjoint_b=True)
  
      # TODO(tfmot): note that it does not suffice to just have the initializer
      # to derive the shape from, in the case of a constant initializer.
      # The unit test fail without providing the shape.
      self.add_training_weight(
          name="u",
          shape=u.shape,
          dtype=u.dtype,
          initializer=tf.keras.initializers.Constant(u),
      )
      self.add_training_weight(
          name="sv",
          shape=sv.shape,
          dtype=sv.dtype,
          initializer=tf.keras.initializers.Constant(sv),
      )
  
file_path: model-optimization/tensorflow_model_optimization/python/core/common/keras/compression/algorithms/same_training_and_inference.py
imports: from typing import List
  import tensorflow
  from tensorflow_model_optimization.python.core.common.keras.compression import algorithm
label_desc: Use image node_id 2 for calling the SVD obj's underlying member method code with example usage: obj.init_training_weights(pretrained_weight) without return types