class_docstr: Defines how to apply the EPR algorithm.
class_docstr_tok: ["Defines","how","to","apply","the","EPR","algorithm","."]
docstr: Initializes training weights needed for reparameterization.
docstr_tok: ["Initializes","training","weights","needed","for","reparameterization","."]
code: def _init_training_weights_reparam(
      self, pretrained_weight: tf.Tensor
  ) -> Tuple[tf.TensorShape, tf.dtypes.DType, str]:
      """Initializes training weights needed for reparameterization."""
      shape = pretrained_weight.shape
      dtype = pretrained_weight.dtype
      weight_name = "bias" if shape.rank == 1 else "kernel"
  
      if 1 <= shape.rank <= 2:
          # Bias or dense kernel.
          self.add_training_weight(
              name=weight_name,
              shape=shape,
              dtype=dtype,
              initializer=tf.keras.initializers.Constant(
                  pretrained_weight
              ),
          )
          prior_shape = tf.TensorShape(())
      elif 3 <= shape.rank <= 4:
          # Convolution kernel.
          kernel_shape = tf.shape(pretrained_weight)
          if shape.rank == 3:
              kernel_rdft = tf.signal.rfft(
                  tf.transpose(pretrained_weight, (1, 2, 0))
              )
          else:
              kernel_rdft = tf.signal.rfft2d(
                  tf.transpose(pretrained_weight, (2, 3, 0, 1))
              )
          kernel_rdft = tf.bitcast(kernel_rdft, tf.float32)
          kernel_rdft /= tf.sqrt(
              tf.cast(tf.reduce_prod(kernel_shape[:-2]), dtype)
          )
          self.add_training_weight(
              name="kernel_rdft",
              shape=kernel_rdft.shape,
              dtype=kernel_rdft.dtype,
              initializer=tf.keras.initializers.Constant(kernel_rdft),
          )
          self.add_training_weight(
              name="kernel_shape",
              shape=kernel_shape.shape,
              dtype=kernel_shape.dtype,
              # TODO(jballe): If False, breaks optimize.create_layer_for_training().
              # If True, throws warnings that int tensors have no gradient.
              # trainable=False,
              initializer=tf.keras.initializers.Constant(kernel_shape),
          )
          prior_shape = kernel_rdft.shape[2:]
      else:
          raise ValueError(
              f"Expected bias or kernel tensor with rank between 1 and 4, received "
              f"shape {shape}."
          )
  
      # Logarithm of quantization step size.
      log_step = tf.fill(prior_shape, tf.constant(-4, dtype=dtype))
      self.add_training_weight(
          name=f"{weight_name}_log_step",
          shape=log_step.shape,
          dtype=log_step.dtype,
          initializer=tf.keras.initializers.Constant(log_step),
      )
  
      return prior_shape, dtype, weight_name
  
imports: import functools
  from typing import Callable, List, Tuple
  import tensorflow
  import tensorflow_compression
  from tensorflow_model_optimization.python.core.common.keras.compression import algorithm
label_desc: Use image node_id 3 for calling the EPRBase obj's underlying member method code with example usage: obj._init_training_weights_reparam(pretrained_weight) and returns: prior_shape,  dtype,  weight_name