code: def adamax(
      loss_or_grads=None,
      params=None,
      learning_rate=0.002,
      beta1=0.9,
      beta2=0.999,
      epsilon=1e-8,
  ):
      """Adamax updates
  
      Adamax updates implemented as in [1]_. This is a variant of the Adam
      algorithm based on the infinity norm.
  
      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 weighted infinity norm estimates.
      epsilon: float
          Constant for numerical stability.
  
      Returns
      -------
      OrderedDict
          A dictionary mapping each parameter to its update expression
  
      Notes
      -----
      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 = adamax(b, [a], learning_rate=.01)
      >>> isinstance(updates, dict)
      True
      >>> optimizer = adamax(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(adamax, **_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 / (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,
          )
          u_prev = pytensor.shared(
              np.zeros(value.shape, dtype=value.dtype),
              shape=param.type.shape,
          )
  
          m_t = beta1 * m_prev + (one - beta1) * g_t
          u_t = pt.maximum(beta2 * u_prev, abs(g_t))
          step = a_t * m_t / (u_t + epsilon)
  
          updates[m_prev] = m_t
          updates[u_prev] = u_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 13 for calling a global function with example usage: adamax(loss_or_grads, params, learning_rate, beta1, beta2, epsilon) and returns: updates, partial