code: def shorrocks_index(A):
      r
  
      A = np.asarray(A)  # Convert to array if not already
      m, n = A.shape
  
      if m != n:
          raise ValueError('A must be a square matrix')
  
      diag_sum = np.diag(A).sum()
  
      return (m - diag_sum) / (m - 1)
docs: r"""
      Implements Shorrocks mobility index
  
      Parameters
      -----------
      A : array_like(float)
          Square matrix with transition probabilities (mobility matrix) of
          dimension m
  
      Returns
      --------
      Shorrocks index: float
          The Shorrocks mobility index calculated as
  
          .. math::
              
              s(A) = \frac{m - \sum_j a_{jj} }{m - 1} \in (0, 1)
  
          An index equal to 0 indicates complete immobility.
  
      References
      -----------
      .. [1] Wealth distribution and social mobility in the US: A quantitative approach
         (Benhabib, Bisin, Luo, 2017).
         https://www.econ.nyu.edu/user/bisina/RevisionAugust.pdf