text: import geoopt
  import torch
  import numpy as np
  import pytest
  
  
  @pytest.mark.parametrize("params", [dict(lr=1e-2), dict(lr=1, amsgrad=True)])
  def test_adam_lorentz(params):
      lorentz = geoopt.manifolds.Lorentz(k=torch.Tensor([1.0]))
      torch.manual_seed(42)
      with torch.no_grad():
          X = geoopt.ManifoldParameter(torch.randn(20, 10), manifold=lorentz).proj_()
      Xstar = torch.randn(20, 10)
      Xstar.set_(lorentz.projx(Xstar))
  
      def closure():
          optim.zero_grad()
          loss = (Xstar - X).pow(2).sum()
          loss.backward()
          return loss.item()
  
      optim = geoopt.optim.RiemannianAdam([X], stabilize=4500, **params)
      assert optim.param_groups[0]["stabilize"] == 4500
      for _ in range(10000):
          if (Xstar - X).norm() < 1e-5:
              break
          optim.step(closure)
      assert X.is_contiguous()
      np.testing.assert_allclose(X.data, Xstar, atol=1e-5, rtol=1e-5)
      optim.load_state_dict(optim.state_dict())
      optim.step(closure)
  
  
  @pytest.mark.parametrize("params", [dict(lr=1e-2), dict(lr=1, amsgrad=True)])
  def test_adam_stiefel(params):
      stiefel = geoopt.manifolds.Stiefel()
      torch.manual_seed(42)
      with torch.no_grad():
          X = geoopt.ManifoldParameter(torch.randn(20, 10), manifold=stiefel).proj_()
      Xstar = torch.randn(20, 10)
      Xstar.set_(stiefel.projx(Xstar))
  
      def closure():
          optim.zero_grad()
          loss = (X - Xstar).pow(2).sum()
          # manifold constraint that makes optimization hard if violated
          loss += (X.t() @ X - torch.eye(X.shape[1])).pow(2).sum() * 100
          loss.backward()
          return loss.item()
  
      optim = geoopt.optim.RiemannianAdam([dict(params=[X], stabilize=4500)], **params)
      assert optim.param_groups[0]["stabilize"] == 4500
      assert (X - Xstar).norm() > 1e-5
      for _ in range(10000):
          if (X - Xstar).norm() < 1e-5:
              break
          optim.step(closure)
      assert X.is_contiguous()
      np.testing.assert_allclose(X.data, Xstar, atol=1e-5, rtol=1e-5)
      optim.load_state_dict(optim.state_dict())
      optim.step(closure)
  
  
  def test_adam_poincare():
      torch.manual_seed(44)
      manifold = geoopt.PoincareBall()
      ideal = torch.tensor([0.5, 0.5])
      start = torch.randn(2) / 2
      start = manifold.expmap0(start)
      start = geoopt.ManifoldParameter(start, manifold=manifold)
  
      def closure():
          optim.zero_grad()
          loss = manifold.dist(start, ideal) ** 2
          loss.backward()
          return loss.item()
  
      optim = geoopt.optim.RiemannianAdam([start], lr=1e-2)
  
      for _ in range(2000):
          optim.step(closure)
      np.testing.assert_allclose(start.data, ideal, atol=1e-5, rtol=1e-5)
  
  
  @pytest.mark.parametrize("params", [dict(lr=1e-2), dict(lr=1, amsgrad=True)])
  def test_adam_StereographicProductManifold(params):
      torch.manual_seed(42)
      pman = geoopt.StereographicProductManifold(
          (geoopt.PoincareBall(), 2),
          (geoopt.SphereProjection(), 2),
          (geoopt.Stereographic(k=0), 2),
      )
      ideal = pman.random(6)
      start = geoopt.ManifoldParameter(pman.random(6), manifold=pman)
  
      def closure():
          optim.zero_grad()
          loss = pman.dist(start, ideal) ** 2
          loss.backward()
          return loss.item()
  
      optim = geoopt.optim.RiemannianAdam([start], **params)
  
      for _ in range(2000):
          optim.step(closure)
      np.testing.assert_allclose(start.data, ideal, atol=1e-5, rtol=1e-5)
  