code: def test_kmax(model_test, data):
      "Test out the k-max semiring"
      K = 2
      gen = Gen(model_test, data, KMaxSemiring(K))
      max1 = gen.model(MaxSemiring).sum(gen.vals)
      alpha = gen.struct.sum(gen.vals, _raw=True)
  
      # 2max is less than max.
      assert (alpha[0] == max1).all()
      assert (alpha[1] <= max1).all()
  
      topk = gen.struct.marginals(gen.vals, _raw=True)
      argmax = gen.model(MaxSemiring).marginals(gen.vals)
  
      # Argmax is different than 2-argmax
      assert (topk[0] == argmax).all()
      assert (topk[1] != topk[0]).any()
  
      if model_test != "Dep":
          log_probs = gen.enum(MaxSemiring)[1]
          tops = torch.topk(torch.cat(log_probs, dim=0), 5, 0)[0]
          assert torch.isclose(
              gen.struct.score(topk[1], gen.vals), alpha[1]
          ).all()
          for k in range(K):
              assert (torch.isclose(alpha[k], tops[k])).all()
  
imports: from torch_struct import CKY, CKY_CRF, DepTree, LinearChain, SemiMarkov, Alignment, deptree_nonproj, deptree_part
  from torch_struct import LogSemiring, CheckpointSemiring, CheckpointShardSemiring, KMaxSemiring, SparseMaxSemiring, MaxSemiring, StdSemiring, EntropySemiring
  from .extensions import LinearChainTest, SemiMarkovTest, DepTreeTest, CKYTest, CKY_CRFTest, test_lookup
  import torch
  from hypothesis import given
  from hypothesis.strategies import integers, data, sampled_from
  import pytest
  from hypothesis import settings
label_desc: Use image node_id 5 for calling a global function with example usage: test_kmax(model_test, data) without return types