Enhance Docstrings in aepsych/kernels (facebookresearch#427)

Summary:
One of several PRs addressing issue facebookresearch#366 to improve docstring coverage.

Improves documentation in `aepsych/kernels` for better clarity and consistency.
- Adds missing docstrings to functions and methods across kernel modules.
- Updates existing docstrings with refined type hints and a unified structure.

Pull Request resolved: facebookresearch#427

Reviewed By: crasanders

Differential Revision: D65827615

Pulled By: JasonKChow

fbshipit-source-id: 76673b5810cc2560853c8a6d6767d0baade694a2

Author: yalsaffar

aepsych/kernels/pairwisekernel.py

@@ -16,8 +16,13 @@ class PairwiseKernel(Kernel):
     """
 
     def __init__(
-        self, latent_kernel: Any, is_partial_obs: bool = False, **kwargs
+          self, latent_kernel: Kernel, is_partial_obs: bool=False, **kwargs
     ) -> None:
+        """
+         Args:
+            latent_kernel (Kernel): The underlying kernel used to compute the covariance for the GP.
+            is_partial_obs (bool): If the kernel should handle partial observations. Defaults to False.
+        """
         super(PairwiseKernel, self).__init__(**kwargs)
 
         self.latent_kernel = latent_kernel
@@ -30,20 +35,20 @@ def forward(
         TODO: make last_batch_dim work properly
 
         d must be 2*k for integer k, k is the dimension of the latent space
+
         Args:
-            :attr:`x1` (Tensor `n x d` or `b x n x d`):
-                First set of data
-            :attr:`x2` (Tensor `m x d` or `b x m x d`):
-                Second set of data
-            :attr:`diag` (bool):
-                Should the Kernel compute the whole kernel, or just the diag?
+            x1 (torch.Tensor): A `b x n x d` or `n x d` tensor, where `d = 2k` and `k` is the dimension of the latent space.
+            x2 (torch.Tensor): A `b x m x d` or `m x d` tensor, where `d = 2k` and `k` is the dimension of the latent space.
+            diag (bool): Should the Kernel compute the whole covariance matrix or just the diagonal? Defaults to False.
+            
 
         Returns:
-            :class:`Tensor` or :class:`gpytorch.lazy.LazyTensor`.
-                The exact size depends on the kernel's evaluation mode:
+            torch.Tensor (or :class:`gpytorch.lazy.LazyTensor`) : A `b x n x m` or `n x m` tensor representing
+            the covariance matrix between `x1` and `x2`. 
+            The exact size depends on the kernel's evaluation mode:
+            * `full_covar`: `n x m` or `b x n x m`
+            * `diag`: `n` or `b x n`
 
-                * `full_covar`: `n x m` or `b x n x m`
-                * `diag`: `n` or `b x n`
         """
         if self.is_partial_obs:
             d: Union[torch.Tensor, int] = x1.shape[-1] - 1

aepsych/kernels/rbf_partial_grad.py

@@ -30,6 +30,21 @@ class RBFKernelPartialObsGrad(RBFKernelGrad):
     def forward(
         self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **params: Any
     ) -> torch.Tensor:
+        """Computes the covariance matrix between x1 and x2 based on the RBF
+        
+        Args:
+            x1 (torch.Tensor): A `b x n x d` or `n x d` tensor, where `d = 2k` and `k` is the dimension of the latent space.
+            x2 (torch.Tensor): A `b x m x d` or `m x d` tensor, where `d = 2k` and `k` is the dimension of the latent space.
+            diag (bool): Should the Kernel compute the whole covariance matrix (False) or just the diagonal (True)? Defaults to False.
+
+        
+        
+        Returns:
+            torch.Tensor: A `b x n x m` or `n x m` tensor representing the covariance matrix between `x1` and `x2`.
+            The exact size depends on the kernel's evaluation mode:
+            * `full_covar`: `n x m` or `b x n x m`
+            * `diag`: `n` or `b x n`
+        """
         # Extract grad index from each
         grad_idx1 = x1[..., -1].to(dtype=torch.long)
         grad_idx2 = x2[..., -1].to(dtype=torch.long)