code: def get_KxX_dx(
      gp: Model, x: Tensor, kernel_type: str = "rbf"
  ) -> Tensor:
      """Computes the analytic derivative of the kernel K(x,X) w.r.t. x.
      Args:
          gp: Botorch model.
          x: (n x D) Test points.
          kernel_type: Takes "rbf" or "matern"
      Returns:
          Tensor (n x D) The derivative of the kernel K(x,X) w.r.t. x.
      """
      X = gp.train_inputs[0]
      D = X.shape[1]
      N = X.shape[0]
      n = x.shape[0]
      lengthscale = gp.covar_module.base_kernel.lengthscale.detach()
      if kernel_type == "rbf":
          K_xX = gp.covar_module(x, X).evaluate()
          part1 = (
              -torch.eye(D, device=x.device, dtype=x.dtype)
              / lengthscale**2
          )
          part2 = x.view(n, 1, D) - X.view(1, N, D)
          return part1 @ (part2 * K_xX.view(n, N, 1)).transpose(1, 2)
      # Else, we have a Matern kernel
      mean = x.reshape(-1, x.size(-1)).mean(0)[(None,) * (x.dim() - 1)]
      x1_ = (x - mean).div(lengthscale)
      x2_ = (X - mean).div(lengthscale)
      distance = gp.covar_module.covar_dist(x1_, x2_)
      exp_component = torch.exp(
          -math.sqrt(5.0) * distance
      )  # pyre-ignore
      constant_component = (-5.0 / 3.0) * distance - (
          5.0 * math.sqrt(5.0) / 3.0
      ) * (distance**2)
      sigma_f = gp.covar_module.outputscale.detach()
      part1 = torch.eye(D, device=lengthscale.device) / lengthscale
      part2 = (
          x1_.view(n, 1, D) - x2_.view(1, N, D)
      ) / distance.unsqueeze(2)
      total_k = sigma_f * constant_component * exp_component
      total = part1 @ (part2 * total_k.view(n, N, 1)).transpose(1, 2)
      return total
  
file_path: Ax/ax/utils/sensitivity/derivative_gp.py
imports: import math
  import torch
  from botorch.models.model import Model
  from gpytorch.distributions import MultivariateNormal
  from torch import Tensor
label_desc: Use image node_id 2 for calling a global function with example usage: get_KxX_dx(gp, x, kernel_type) and returns: total, unknown