docstr: Tests SVD with rectangular input.
docstr_tok: ["Tests","SVD","with","rectangular","input","."]
code: def testSvdWithRectangularInput(self, m, n, log_cond, full_matrices):
      """Tests SVD with rectangular input."""
      with jax.default_matmul_precision("float32"):
          a = np.random.uniform(low=0.3, high=0.9, size=(m, n)).astype(
              _SVD_TEST_DTYPE
          )
          u, s, v = osp_linalg.svd(a, full_matrices=False)
          cond = 10**log_cond
          s = jnp.linspace(cond, 1, min(m, n))
          a = (u * s) @ v
          a = a.astype(complex) * (1 + 1j)
  
          osp_linalg_fn = functools.partial(
              osp_linalg.svd, full_matrices=full_matrices
          )
          actual_u, actual_s, actual_v = svd.svd(
              a, full_matrices=full_matrices
          )
  
          k = min(m, n)
          if m > n:
              unitary_u = jnp.real(actual_u.T.conj() @ actual_u)
              unitary_v = jnp.real(actual_v.T.conj() @ actual_v)
              unitary_u_size = m if full_matrices else k
              unitary_v_size = k
          else:
              unitary_u = jnp.real(actual_u @ actual_u.T.conj())
              unitary_v = jnp.real(actual_v @ actual_v.T.conj())
              unitary_u_size = k
              unitary_v_size = n if full_matrices else k
  
          _, expected_s, _ = osp_linalg_fn(a)
  
          svd_fn = lambda a: svd.svd(a, full_matrices=full_matrices)
          args_maker = lambda: [a]
  
          with self.subTest("Test JIT compatibility"):
              self._CompileAndCheck(svd_fn, args_maker)
  
          with self.subTest("Test unitary u."):
              self.assertAllClose(
                  np.eye(unitary_u_size),
                  unitary_u,
                  rtol=_SVD_RTOL,
                  atol=2e-3,
              )
  
          with self.subTest("Test unitary v."):
              self.assertAllClose(
                  np.eye(unitary_v_size),
                  unitary_v,
                  rtol=_SVD_RTOL,
                  atol=2e-3,
              )
  
          with self.subTest("Test s."):
              self.assertAllClose(
                  expected_s,
                  jnp.real(actual_s),
                  rtol=_SVD_RTOL,
                  atol=1e-6,
              )
  
file_path: jax/tests/svd_test.py
imports: import functools
  import jax
  import jax.numpy
  import numpy
  import scipy.linalg
  from jax._src import config
  from jax._src import test_util
  from jax._src.lax import svd
  from absl.testing import absltest
label_desc: Use image node_id 2 for calling the SvdTest obj's underlying member method code with example usage: obj.testSvdWithRectangularInput(m, n, log_cond, full_matrices) without return types