def METHOD_NAME(self, weights_1d):
    """Sets the weights with a 1D tf.Tensor.
    Args:
        weights_1d: a 1D tf.Tensor representing the trainable variables.
    """
    weights = tf.dynamic_partition(weights_1d, self.partitions, self.n_tensors)
    for i, (shape, param) in enumerate(zip(self.shapes, weights)):
        self.trainable_variables[i].assign(tf.reshape(param, shape))