code: def momentum(
      loss_or_grads=None, params=None, learning_rate=1e-3, momentum=0.9
  ):
      """Stochastic Gradient Descent (SGD) updates with momentum
  
      Generates update expressions of the form:
  
      * ``velocity := momentum * velocity - learning_rate * gradient``
      * ``param := param + velocity``
  
      Parameters
      ----------
      loss_or_grads: symbolic expression or list of expressions
          A scalar loss expression, or a list of gradient expressions
      params: list of shared variables
          The variables to generate update expressions for
      learning_rate: float or symbolic scalar
          The learning rate controlling the size of update steps
      momentum: float or symbolic scalar, optional
          The amount of momentum to apply. Higher momentum results in
          smoothing over more update steps. Defaults to 0.9.
  
      Returns
      -------
      OrderedDict
          A dictionary mapping each parameter to its update expression
  
      Notes
      -----
      Higher momentum also results in larger update steps. To counter that,
      you can optionally scale your learning rate by `1 - momentum`.
  
      Optimizer can be called without both loss_or_grads and params
      in that case partial function is returned
  
      See Also
      --------
      apply_momentum: Generic function applying momentum to updates
      nesterov_momentum: Nesterov's variant of SGD with momentum
  
      Examples
      --------
      >>> a = pytensor.shared(1.)
      >>> b = a*2
      >>> updates = momentum(b, [a], learning_rate=.01)
      >>> isinstance(updates, dict)
      True
      >>> optimizer = momentum(learning_rate=.01)
      >>> callable(optimizer)
      True
      >>> updates = optimizer(b, [a])
      >>> isinstance(updates, dict)
      True
      """
      if loss_or_grads is None and params is None:
          return partial(
              pm.updates.momentum, **_get_call_kwargs(locals())
          )
      elif loss_or_grads is None or params is None:
          raise ValueError(
              "Please provide both `loss_or_grads` and `params` to get updates"
          )
      updates = sgd(loss_or_grads, params, learning_rate)
      return apply_momentum(updates, momentum=momentum)
  
imports: from collections import OrderedDict
  from functools import partial
  import numpy
  import pytensor
  import pytensor.tensor
  import pymc
label_desc: Use image node_id 5 for calling a global function with example usage: momentum(loss_or_grads, params, learning_rate, momentum) and returns: apply_momentum, partial