def testSchurBatching(self, shape, dtype):
      rng = jtu.rand_default(self.rng())
      batch_size = 10
      shape = (batch_size, ) + shape
      args = rng(shape, dtype)
      reconstruct = vmap(lambda S, T: S @ T @ jnp.conj(S.T))

      Ts, Ss = vmap(lax.linalg.schur)(args)
      self.assertAllClose(reconstruct(Ss, Ts), args, atol=1e-4)