parameters: (self, new_grads, damping_step, gradients)
return_statement: return new_grads
docstring: get second gradients for thor
docstring_summary: get second gradients for thor
docstring_tokens: get second gradients for thor
function: def _get_second_gradients(self, new_grads, damping_step, gradients):
          """get second gradients for thor"""
          params_len = len(self.params)
          if self.conv_layer_count > 0:
              new_grads = self._get_second_gradients_one(params_len, gradients, new_grads)
          else:
              for i in range(params_len):
                  g = gradients[i]
                  thor_layer_count = self.weight_fim_idx_map[i]
                  layer_type = self.weight_layertype_idx_map[i]
                  if layer_type == Embedding:
                      temp_a_ori = self.matrix_a_cov[thor_layer_count]
                      temp_g = self.matrix_g_cov[thor_layer_count]
                      temp_a = self.expand(temp_a_ori, 1)
                      g = self.mul(temp_a, g)
                      temp_g = self.cast(temp_g, mstype.float16)
                      g = self.cast(g, mstype.float16)
                      g = self.matmul(g, temp_g)
                      g = self.cast(g, mstype.float32)
                  elif layer_type == FC:
                      temp_a = self.matrix_a_cov[thor_layer_count]
                      temp_g = self.matrix_g_cov[thor_layer_count]
                      temp_a = self.cast(temp_a, mstype.float16)
                      temp_g = self.cast(temp_g, mstype.float16)
                      g = self.cast(g, mstype.float16)
                      g = self.matmul(temp_g, g)
                      g = self.matmul(g, temp_a)
                      g = self.cast(g, mstype.float32)
                  elif layer_type == LayerNorm:
                      g = self._process_layernorm(damping_step, g)
                  new_grads = new_grads + (g,)
          return new_grads
function_tokens: def _get_second_gradients ( self , new_grads , damping_step , gradients ) : params_len = len ( self . params ) if self . conv_layer_count > 0 : new_grads = self . _get_second_gradients_one ( params_len , gradients , new_grads ) else : for i in range ( params_len ) : g = gradients [ i ] thor_layer_count = self . weight_fim_idx_map [ i ] layer_type = self . weight_layertype_idx_map [ i ] if layer_type == Embedding : temp_a_ori = self . matrix_a_cov [ thor_layer_count ] temp_g = self . matrix_g_cov [ thor_layer_count ] temp_a = self . expand ( temp_a_ori , 1 ) g = self . mul ( temp_a , g ) temp_g = self . cast ( temp_g , mstype . float16 ) g = self . cast ( g , mstype . float16 ) g = self . matmul ( g , temp_g ) g = self . cast ( g , mstype . float32 ) elif layer_type == FC : temp_a = self . matrix_a_cov [ thor_layer_count ] temp_g = self . matrix_g_cov [ thor_layer_count ] temp_a = self . cast ( temp_a , mstype . float16 ) temp_g = self . cast ( temp_g , mstype . float16 ) g = self . cast ( g , mstype . float16 ) g = self . matmul ( temp_g , g ) g = self . matmul ( g , temp_a ) g = self . cast ( g , mstype . float32 ) elif layer_type == LayerNorm : g = self . _process_layernorm ( damping_step , g ) new_grads = new_grads + ( g , ) return new_grads