text: import pytest
  import torch
  
  from piq import KID
  
  
  @pytest.fixture(scope='module')
  def features_target_normal() -> torch.Tensor:
      return torch.rand(1000, 20)
  
  
  @pytest.fixture(scope='module')
  def features_prediction_normal() -> torch.Tensor:
      return torch.rand(1000, 20)
  
  
  @pytest.fixture(scope='module')
  def features_prediction_beta() -> torch.Tensor:
      m = torch.distributions.Beta(torch.FloatTensor([2]), torch.FloatTensor([2]))
      return m.sample([1000, 20]).squeeze()
  
  
  @pytest.fixture(scope='module')
  def features_prediction_constant() -> torch.Tensor:
      return torch.ones(1000, 20)
  
  
  def test_initialization() -> None:
      try:
          KID()
      except Exception as e:
          pytest.fail(f"Unexpected error occurred: {e}")
  
  
  def test_forward(features_target_normal: torch.Tensor, features_prediction_normal: torch.Tensor, ) -> None:
      try:
          metric = KID()
          metric(features_target_normal, features_prediction_normal)
      except Exception as e:
          pytest.fail(f"Unexpected error occurred: {e}")
  
  
  def tes_fails_for_different_dimensions(features_target_normal: torch.Tensor) -> None:
      features_prediction_normal = torch.rand(1000, 21)
      metric = KID()
      with pytest.raises(AssertionError):
          metric(features_target_normal, features_prediction_normal)
  
  
  def test_works_for_different_number_of_images_in_stack(features_target_normal: torch.Tensor) -> None:
      features_prediction_normal = torch.rand(1010, 20)
      metric = KID()
      metric(features_target_normal, features_prediction_normal)
  
  
  def test_returns_variance(features_target_normal: torch.Tensor, features_prediction_normal: torch.Tensor) -> None:
      metric = KID(ret_var=True)
      result = metric(features_target_normal, features_prediction_normal)
      print(result)
      assert len(result) == 2, \
          f'Expected to get score and variance, got {result}'
  