code: def _constant_svd(
      a: Any,
      return_nan: bool,
      full_matrices: bool,
      compute_uv: bool = True,
  ) -> Any | Sequence[Any]:
      """SVD on matrix of all zeros."""
      m, n = a.shape
      k = min(m, n)
      s = jnp.where(
          return_nan,
          jnp.full(shape=(k,), fill_value=jnp.nan, dtype=a.real.dtype),
          jnp.zeros(shape=(k,), dtype=a.real.dtype),
      )
      if compute_uv:
          fill_value = (
              jnp.nan + 1j * jnp.nan
              if jnp.issubdtype(a.dtype, jnp.complexfloating)
              else jnp.nan
          )
          if full_matrices:
              u = jnp.where(
                  return_nan,
                  jnp.full((m, m), fill_value, dtype=a.dtype),
                  jnp.eye(m, m, dtype=a.dtype),
              )
              vh = jnp.where(
                  return_nan,
                  jnp.full((n, n), fill_value, dtype=a.dtype),
                  jnp.eye(n, n, dtype=a.dtype),
              )
          else:
              u = jnp.where(
                  return_nan,
                  jnp.full((m, k), fill_value, dtype=a.dtype),
                  jnp.eye(m, k, dtype=a.dtype),
              )
              vh = jnp.where(
                  return_nan,
                  jnp.full((k, n), fill_value, dtype=a.dtype),
                  jnp.eye(k, n, dtype=a.dtype),
              )
          return (u, s, vh)
      else:
          return s
  
file_path: jax/jax/_src/lax/svd.py
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  from typing import Any
  import jax
  import jax.numpy
  from jax import lax
  from jax._src import core
label_desc: Use image node_id 1 for calling a global function with example usage: _constant_svd(a, return_nan, full_matrices, compute_uv) and returns: u,  s,  vh, s