code: def diag(A, k=0):
      
      if isinstance(A, Poly):
          core, core_new = A.A, {}
          for key in A.keys:
              core_new[key] = numpy.diag(core[key], k)
  
          return Poly(core_new, A.dim, None, A.dtype)
  
      return numpy.diag(A, k)
docs: Extract or construct a diagonal polynomial array.