code: def _parallel_update_step(
      self,
      grads,
      variables,
      learning_rate,
  ):
      keras_variables = variables
      variables = [v.value for v in variables]
  
      dtype = variables[0].dtype
      lr = ops.cast(learning_rate, dtype)
  
      local_step = ops.cast(self.iterations + 1, dtype)
      next_step = ops.cast(self.iterations + 2, dtype)
      decay = ops.cast(0.96, dtype)
      beta_1 = ops.cast(self.beta_1, dtype)
      beta_2 = ops.cast(self.beta_2, dtype)
      u_t = beta_1 * (1.0 - 0.5 * (ops.power(decay, local_step)))
      u_t_1 = beta_1 * (1.0 - 0.5 * (ops.power(decay, next_step)))
      u_product_t = self._u_product.value * u_t
      u_product_t_1 = u_product_t * u_t_1
      beta_2_power = ops.power(beta_2, local_step)
  
      self._u_product.assign(u_product_t)
  
      m_list = [
          self._momentums[self._get_variable_index(variable)].value
          for variable in keras_variables
      ]
      v_list = [
          self._velocities[self._get_variable_index(variable)].value
          for variable in keras_variables
      ]
  
      torch._foreach_mul_(m_list, self.beta_1)
      torch._foreach_add_(m_list, grads, alpha=1 - self.beta_1)
  
      torch._foreach_mul_(v_list, self.beta_2)
      torch._foreach_add_(
          v_list,
          torch._foreach_mul(grads, grads),
          alpha=1 - self.beta_2,
      )
  
      m_hat_list = torch._foreach_add(
          torch._foreach_div(
              torch._foreach_mul(m_list, u_t_1),
              1 - core.convert_to_numpy(u_product_t_1),
          ),
          torch._foreach_div(
              torch._foreach_mul(grads, 1 - u_t),
              1 - core.convert_to_numpy(u_product_t),
          ),
      )
  
      v_hat_list = torch._foreach_div(v_list, 1 - beta_2_power)
  
      torch._foreach_add_(
          variables,
          torch._foreach_div(
              torch._foreach_mul(m_hat_list, lr),
              torch._foreach_add(
                  torch._foreach_sqrt(v_hat_list), self.epsilon
              ),
          ),
          alpha=-1,
      )
  
file_path: keras/keras/backend/torch/optimizers/torch_nadam.py
imports: import torch
  from keras import ops
  from keras import optimizers
  from keras.backend.torch import core
  from keras.backend.torch.optimizers import torch_parallel_optimizer
label_desc: Use image node_id 1 for calling the Nadam obj's underlying member method code with example usage: obj._parallel_update_step(grads, variables, learning_rate) without return types