code: def adam(
      loss_or_grads=None,
      params=None,
      learning_rate=0.001,
      beta1=0.9,
      beta2=0.999,
      epsilon=1e-8,
  ):
      """Adam updates
  
      Adam updates implemented as in [1]_.
  
      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
          Learning rate
      beta1: float
          Exponential decay rate for the first moment estimates.
      beta2: float
          Exponential decay rate for the second moment estimates.
      epsilon: float
          Constant for numerical stability.
  
      Returns
      -------
      OrderedDict
          A dictionary mapping each parameter to its update expression
  
      Notes
      -----
      The paper [1]_ includes an additional hyperparameter lambda. This is only
      needed to prove convergence of the algorithm and has no practical use
      (personal communication with the authors), it is therefore omitted here.
  
      Optimizer can be called without both loss_or_grads and params
      in that case partial function is returned
  
      References
      ----------
      .. [1] Kingma, Diederik, and Jimmy Ba (2014):
             Adam: A Method for Stochastic Optimization.
             arXiv preprint arXiv:1412.6980.
  
      Examples
      --------
      >>> a = pytensor.shared(1.)
      >>> b = a*2
      >>> updates = adam(b, [a], learning_rate=.01)
      >>> isinstance(updates, dict)
      True
      >>> optimizer = adam(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(adam, **_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"
          )
      all_grads = get_or_compute_grads(loss_or_grads, params)
      t_prev = pytensor.shared(pm.pytensorf.floatX(0.0))
      updates = OrderedDict()
  
      # Using pytensor constant to prevent upcasting of float32
      one = pt.constant(1)
  
      t = t_prev + 1
      a_t = (
          learning_rate * pt.sqrt(one - beta2**t) / (one - beta1**t)
      )
  
      for param, g_t in zip(params, all_grads):
          value = param.get_value(borrow=True)
          m_prev = pytensor.shared(
              np.zeros(value.shape, dtype=value.dtype),
              shape=param.type.shape,
          )
          v_prev = pytensor.shared(
              np.zeros(value.shape, dtype=value.dtype),
              shape=param.type.shape,
          )
  
          m_t = beta1 * m_prev + (one - beta1) * g_t
          v_t = beta2 * v_prev + (one - beta2) * g_t**2
          step = a_t * m_t / (pt.sqrt(v_t) + epsilon)
  
          updates[m_prev] = m_t
          updates[v_prev] = v_t
          updates[param] = param - step
  
      updates[t_prev] = t
      return updates
  
imports: from collections import OrderedDict
  from functools import partial
  import numpy
  import pytensor
  import pytensor.tensor
  import pymc
label_desc: Use image node_id 12 for calling a global function with example usage: adam(loss_or_grads, params, learning_rate, beta1, beta2, epsilon) and returns: updates, partial