code: def _scale_by_learning_rate(
      learning_rate: optax.ScalarOrSchedule,
      flip_sign: bool = True,
  ) -> optax.GradientTransformation:
      m = -1 if flip_sign else 1
      if callable(learning_rate):
          return optax.scale_by_schedule(
              lambda count: m * learning_rate(count)
          )
      return optax.scale(m * learning_rate)
  
imports: from collections.abc import Callable
  from typing import Any
  from typing import Optional
  from typing import Union
  import optax
label_desc: Use image node_id 1 for calling a global function with example usage: _scale_by_learning_rate(learning_rate, flip_sign) and returns: optax, optax