//  [python]
def optimize_loss(
    loss,
    learning_rate,
    optimizer_name,
    clip_gradients,
    lr_schedule,
    lr_warmup,
    total_num_steps,
    summarize_grads,
    mixed_precision,
    b1,
    b2,
    epsilon,
    l2_reg,
    vector_l2,
    accumulate_steps,
    max_training_hours=None,
    colocate_gradients_with_ops=True,
    acc_grads_on_cpu=True,
):
    global_step = tf.compat.v1.train.get_or_create_global_step()

    with tf.compat.v1.variable_scope("OptimizeLoss", [loss, global_step]):
        tf.compat.v1.summary.scalar("loss", loss)

        training_fraction = tf.cast(global_step, dtype=tf.float32) / total_num_steps
        if max_training_hours is not None:
            start_time = tf.compat.v1.get_variable(
                initializer=tf.timestamp,
                trainable=False,
                name="start_time",
                dtype=tf.float32,
            )
            training_fraction = tf.minimum(
                tf.maximum(
                    training_fraction,
                    tf.cast(
                        (tf.timestamp() - start_time) / (max_training_hours * 3600),
                        tf.float32,
                    ),
                ),
                1.0,
            )
        training_fraction_var = tf.compat.v1.get_variable(
            initializer=-1.0,
            trainable=False,
            name="training_fraction",
            dtype=tf.float32,
        )
        learning_rate = tf.maximum(
            0.0,
            learning_rate * schedules[lr_schedule](training_fraction, warmup=lr_warmup),
        )
        tf.compat.v1.summary.scalar("learning_rate", learning_rate)

        opt = get_optimizer(
            optimizer_name=optimizer_name,
            learning_rate=learning_rate,
            b1=b1,
            b2=b2,
            epsilon=epsilon,
            l2_reg=l2_reg,
            vector_l2=vector_l2,
            accumulate_steps=accumulate_steps,
            mixed_precision=mixed_precision,
            acc_grads_on_cpu=acc_grads_on_cpu,
        )
        variables = tf.compat.v1.trainable_variables()
        variables = [v for v in variables if is_numerical_tensor(v)]

        # Compute gradients.
        gradients = list(
            zip(tf.gradients(ys=loss, xs=variables, name="gradients"), variables)
        )
        for g, v in gradients:
            if g is None:
                LOGGER.warning("Variable {} has None grads".format(v.name))

        tf.compat.v1.summary.scalar(
            "global_norm/gradient_norm",
            tf.linalg.global_norm([g[0] for g in gradients]),
        )

        gradients = _clip_gradients_by_norm(gradients, clip_gradients)
        tf.compat.v1.summary.scalar(
            "global_norm/clipped_gradient_norm",
            tf.linalg.global_norm(list(zip(*gradients))[0]),
        )

        # Add histograms for variables, gradients and gradient norms.
        if summarize_grads:
            for gradient, variable in gradients:
                if isinstance(gradient, tf.IndexedSlices):
                    grad_values = gradient.values
                else:
                    grad_values = gradient

                if grad_values is not None:
                    var_name = variable.name.replace(":", "_")
                    tf.compat.v1.summary.histogram(
                        "gradients/%s" % var_name, grad_values
                    )
                    tf.compat.v1.summary.scalar(
                        "gradient_norm/%s" % var_name,
                        tf.linalg.global_norm([grad_values]),
                    )

        # Create gradient updates.
        with tf.compat.v1.control_dependencies(
            [global_step.assign_add(1), training_fraction_var.assign(training_fraction)]
        ):
            grad_updates = opt.apply_gradients(gradients, name="train")

    return grad_updates