class_docstr: EXPERIMENTAL Discrete Cosine Transform-augmented
  :class:`~pyro.optim.clipped_adam.ClippedAdam` optimizer.
  
  This acts like :class:`~pyro.optim.clipped_adam.ClippedAdam` on most
  parameters, but if a parameter has an attribute ``._pyro_dct_dim``
  indicating a time dimension, this creates a secondary optimize in the
  frequency domain. This is useful for parameters of time series models.
  
  :param params: iterable of parameters to optimize or dicts defining parameter groups
  :param float lr: learning rate (default: 1e-3)
  :param tuple betas: coefficients used for computing
      running averages of gradient and its square (default: (0.9, 0.999))
  :param float eps: term added to the denominator to improve
      numerical stability (default: 1e-8)
  :param float clip_norm: magnitude of norm to which gradients are clipped (default: 10.0)
  :param float lrd: rate at which learning rate decays (default: 1.0)
  :param bool subsample_aware: whether to update gradient statistics only for
      those elements that appear in a subsample (default: False).
class_docstr_tok: ["EXPERIMENTAL","Discrete","Cosine","Transform-augmented",":","class",":","`","~pyro.optim.clipped_adam.ClippedAdam","`","optimizer",".","This","acts","like",":","class",":","`","~pyro.optim.clipped_adam.ClippedAdam","`","on","most","parameters",",","but","if","a","parameter","has","an","attribute","``","._pyro_dct_dim","``","indicating","a","time","dimension",",","this","creates","a","secondary","optimize","in","the","frequency","domain",".","This","is","useful","for","parameters","of","time","series","models",".",":","param","params",":","iterable","of","parameters","to","optimize","or","dicts","defining","parameter","groups",":","param","float","lr",":","learning","rate","(","default",":","1e-3",")",":","param","tuple","betas",":","coefficients","used","for","computing","running","averages","of","gradient","and","its","square","(","default",":","(","0.9",",","0.999",")",")",":","param","float","eps",":","term","added","to","the","denominator","to","improve","numerical","stability","(","default",":","1e-8",")",":","param","float","clip_norm",":","magnitude","of","norm","to","which","gradients","are","clipped","(","default",":","10.0",")",":","param","float","lrd",":","rate","at","which","learning","rate","decays","(","default",":","1.0",")",":","param","bool","subsample_aware",":","whether","to","update","gradient","statistics","only","for","those","elements","that","appear","in","a","subsample","(","default",":","False",")","."]
code: def _step_param_subsample(self, group: Dict, p, subsample) -> None:
      mask = _get_mask(p, subsample)
  
      grad = p.grad.data.masked_select(mask)
      grad.clamp_(-group["clip_norm"], group["clip_norm"])
  
      # Transform selected parameters via dct.
      time_dim = getattr(p, "_pyro_dct_dim", None)
      if time_dim is not None:
          duration = p.size(time_dim)
          grad = _transform_forward(grad, time_dim, duration)
  
      state = self.state[p]
  
      # State initialization
      if len(state) == 0:
          state["step"] = torch.zeros_like(p)
          # Exponential moving average of gradient values
          state["exp_avg"] = torch.zeros_like(p)
          # Exponential moving average of squared gradient values
          state["exp_avg_sq"] = torch.zeros_like(p)
  
      beta1, beta2 = group["betas"]
  
      state_step = state["step"].masked_select(mask).add_(1)
      state["step"].masked_scatter_(mask, state_step)
  
      # Decay the first and second moment running average coefficient
      exp_avg = (
          state["exp_avg"]
          .masked_select(mask)
          .mul_(beta1)
          .add_(grad, alpha=1 - beta1)
      )
      state["exp_avg"].masked_scatter_(mask, exp_avg)
  
      exp_avg_sq = (
          state["exp_avg_sq"]
          .masked_select(mask)
          .mul_(beta2)
          .addcmul_(grad, grad, value=1 - beta2)
      )
      state["exp_avg_sq"].masked_scatter_(mask, exp_avg_sq)
  
      denom = exp_avg_sq.sqrt_().add_(group["eps"])
  
      bias_correction1 = 1 - beta1**state_step
      bias_correction2 = 1 - beta2**state_step
      step_size = (
          bias_correction2.sqrt_()
          .div_(bias_correction1)
          .mul_(group["lr"])
      )
  
      step = exp_avg.div_(denom)
      if time_dim is not None:
          step = _transform_inverse(step, time_dim, duration)
      p_data = p.data.masked_select(mask).sub_(step.mul_(step_size))
      p.data.masked_scatter_(mask, p_data)
  
file_path: pyro/pyro/optim/dct_adam.py
imports: import math
  from typing import Callable, Dict, Optional, Tuple
  import torch
  from torch.nn.functional import pad
  from torch.optim.optimizer import Optimizer
  from pyro.ops.tensor_utils import dct, idct, next_fast_len
label_desc: Use image node_id 4 for calling the DCTAdam obj's underlying member method code with example usage: obj._step_param_subsample(group, p, subsample) without return types