code: def backward(ctx, output_grads):
      (
          use_biases_t,
          heads_t,
          scale_t,
          matmul2_results,
          dropout_results,
          softmax_results,
          input_lin_results,
          inputs,
          input_weights,
          output_weights,
          dropout_mask,
          dropout_prob_t,
      ) = ctx.saved_tensors
  
      head_dim = inputs.size(2) // heads_t[0]
  
      # Slice out q,k,v from one big Input Linear outuput (should only impact meta data, no copies!)
      # Sequences and heads are combined to make the batch of the Batched GEMM
      # input_lin_results: [seql_q, seqs, heads(16), 3, head_dim(64)]
      # input_lin_results: [seql_q, batches=seqs*heads, 3, head_dim]
      input_lin_results = input_lin_results.view(
          inputs.size(0), inputs.size(1) * heads_t[0], 3, head_dim
      )
      queries = input_lin_results[:, :, 0, :]
      keys = input_lin_results[:, :, 1, :]
      values = input_lin_results[:, :, 2, :]
  
      # Slice out q,k,v from one big set of gradients entering the input linear's bprop  (should only impact meta data, no copies!)
      # The gradients are identical in size to the Input Linear outputs.
      # The tensor is declared before hand to properly slice out query, key, and value grads.
      input_lin_results_grads = torch.empty_like(input_lin_results)
      queries_grads = input_lin_results_grads[:, :, 0, :]
      keys_grads = input_lin_results_grads[:, :, 1, :]
      values_grads = input_lin_results_grads[:, :, 2, :]
  
      # Output Linear GEMM - DGRAD
      # Input1: (data grads)  [seql_q, seqs, embed_dim=heads*head_dim]
      # Input2: (weights)     [ embed_dim, embed_dim ]
      # Output:               [ seql_q, seqs, embed_dim ]
      # GEMM: ( seql_q*seqs x embed_dim ) x ( embed_dim x embed_dim ) = ( seql_q*seqs x embed_dim )
      output_lin_grads = torch.mm(
          output_grads.view(
              output_grads.size(0) * output_grads.size(1),
              output_grads.size(2),
          ),
          output_weights,
      )
      output_lin_grads = output_lin_grads.view(
          output_grads.size(0),
          output_grads.size(1),
          output_weights.size(1),
      )
      # Output Linear GEMM - WGRAD
      # Input1: (data grads)  [seql_q*seqs, embed_dim=heads*head_dim] transpose(0,1)
      # Input2: (activations) [seql_q*seqs, embed_dim ]
      # Output:               [ seql_q, seqs, embed_dim ]
      # GEMM: ( embed_dim x seql_q*seqs ) x ( seql_q*seqs x embed_dim ) = ( embed_dim x embed_dim )
      output_weight_grads = torch.mm(
          output_grads.view(
              output_grads.size(0) * output_grads.size(1),
              output_grads.size(2),
          ).transpose(0, 1),
          matmul2_results.view(
              matmul2_results.size(0) * matmul2_results.size(1),
              matmul2_results.size(2),
          ),
      )
      output_lin_grads = output_lin_grads.view(
          inputs.size(0), inputs.size(1) * heads_t[0], head_dim
      ).transpose(0, 1)
  
      if use_biases_t[0]:
          output_bias_grads = torch.sum(
              output_grads.view(
                  output_grads.size(0) * output_grads.size(1),
                  output_grads.size(2),
              ),
              0,
          )
      else:
          output_bias_grads = None
  
      # Matmul2 - DGRAD1
      # Input1: (data grads)  [seql_q, seqs*heads, head_dim] transpose(0,1)
      # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1).transpose(1,2)
      # Output:               [seqs*heads, seql_q, seql_k]
      # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
      matmul2_dgrad1 = torch.bmm(
          output_lin_grads, values.transpose(0, 1).transpose(1, 2)
      )
      # Matmul2 - DGRAD2
      # Input1: (data grads)  [seql_q, seqs*heads, head_dim] transpose(0,1)
      # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1).transpose(1,2)
      # Output:               [seqs*heads, seql_q, seql_k]
      # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
      values_grads = torch.bmm(
          dropout_results.transpose(1, 2),
          output_lin_grads,
          out=values_grads.transpose(0, 1),
      )
  
      # Mask and Scaling for Dropout (not a publically documented op)
      dropout_grads = torch._masked_scale(
          matmul2_dgrad1, dropout_mask, 1.0 / (1.0 - dropout_prob_t[0])
      )
  
      # Softmax Grad (not a publically documented op)
      softmax_grads = torch._softmax_backward_data(
          dropout_grads, softmax_results, -1, softmax_results.dtype
      )
  
      # Matmul1 - DGRAD1
      # Input1: (data grads)  [seqs*heads, seql_q, seql_k]
      # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1)
      # Output:               [seqs*heads, seql_q, head_dim] transpose(0,1)
      # GEMM: Per batch: ( seql_q x seql_k ) x ( seql_k x head_dim ) = ( seql_q x head_dim )
      queries_grads = torch.baddbmm(
          queries_grads.transpose(0, 1),
          softmax_grads,
          keys.transpose(0, 1),
          out=queries_grads.transpose(0, 1),
          beta=0.0,
          alpha=scale_t[0],
      )
      # Matmul1 - DGRAD2
      # Input1: (data grads)  [seqs*heads, seql_q, seql_k] transpose(1,2)
      # Input2: (activations) [seql_q, seqs*heads, head_dim] transpose(0,1)
      # Output:               [seqs*heads, seql_k, head_dim] transpose(0,1)
      # GEMM: Per batch: ( seql_k x seql_q ) x ( seql_q x head_dim ) = ( seql_k x head_dim )
      keys_grads = torch.baddbmm(
          keys_grads.transpose(0, 1),
          softmax_grads.transpose(1, 2),
          queries.transpose(0, 1),
          out=keys_grads.transpose(0, 1),
          beta=0.0,
          alpha=scale_t[0],
      )
  
      # Input Linear GEMM - DGRAD
      # input1: (data grads) [seql_q, seqs, 3*embed_dim(3072)]
      # input2: (weights)    [embed_dim*3 (3072), embed_dim (1024)]
      # output:              [seql_q, seqs, embed_dim]
      # GEMM: ( (seql_q*seqs) x 3*embed_dim ) x ( 3*embed_dim x embed_dim ) = (seql_q*seqs x embed_dim)
      input_lin_results_grads = input_lin_results_grads.view(
          inputs.size(0) * inputs.size(1), heads_t[0] * 3 * head_dim
      )
      input_grads = torch.mm(input_lin_results_grads, input_weights)
      input_grads = input_grads.view(
          inputs.size(0), inputs.size(1), inputs.size(2)
      )
      # Input Linear GEMM - WGRAD
      # input1: (data grads)  [seql_q*seqs, 3*embed_dim(3072)]
      # input2: (activations) [seql_q*seqs, embed_dim(1024)]
      # output:               [3*embed_dim, embed_dim]
      # GEMM: ( 3*embed_dim x seql_q*seqs ) x ( seql_q*seqs x embed_dim ) = (3*embed_dim x embed_dim)
      input_weight_grads = torch.mm(
          input_lin_results_grads.transpose(0, 1),
          inputs.view(inputs.size(0) * inputs.size(1), inputs.size(2)),
      )
  
      if use_biases_t[0]:
          input_bias_grads = torch.sum(input_lin_results_grads, 0)
      else:
          input_bias_grads = None
  
      return (
          None,
          None,
          None,
          None,
          input_grads,
          input_weight_grads,
          output_weight_grads,
          input_bias_grads,
          output_bias_grads,
          None,
          None,
          None,
      )
  
imports: import torch
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the SelfAttnFunc obj's underlying member method code with example usage: obj.backward(ctx, output_grads) and returns: None,  None,  None,  None,  input_grads,  input_weight_grads,  output_weight_grads,  input_bias_grads,  output_bias_grads,  None,  None,  None