code: def _roots_with_zeros(p: Array, num_leading_zeros: int) -> Array:
      # Avoid lapack errors when p is all zero
      p = _where(len(p) == num_leading_zeros, 1.0, p)
      # Roll any leading zeros to the end & compute the roots
      roots = _roots_no_zeros(roll(p, -num_leading_zeros))
      # Sort zero roots to the end.
      roots = lax.sort_key_val(roots == 0, roots)[1]
      # Set roots associated with num_leading_zeros to NaN
      return _where(
          arange(roots.size) < roots.size - num_leading_zeros,
          roots,
          complex(np.nan, np.nan),
      )
  
imports: from __future__ import annotations
  from functools import partial
  import operator
  import numpy
  from jax import jit
  from jax import lax
  from jax._src import dtypes
  from jax._src import core
  from jax._src.numpy.lax_numpy import arange, argmin, array, asarray, atleast_1d, concatenate, convolve, diag, dot, finfo, full, ones, outer, roll, trim_zeros, trim_zeros_tol, vander, zeros
  from jax._src.numpy.ufuncs import maximum, true_divide, sqrt
  from jax._src.numpy.reductions import all
  from jax._src.numpy import linalg
  from jax._src.numpy.util import check_arraylike, promote_dtypes, promote_dtypes_inexact, _where, _wraps
  from jax._src.typing import Array, ArrayLike
label_desc: Use image node_id 2 for calling a global function with example usage: _roots_with_zeros(p, num_leading_zeros) and returns: _where