code: def test_simple(data, seed):
      model = data.draw(sampled_from([LinearChainCRF]))
      struct = model.struct
      torch.manual_seed(seed)
      vals, (batch, N) = struct._rand()
      lengths = torch.tensor(
          [
              data.draw(integers(min_value=2, max_value=N))
              for b in range(batch - 1)
          ]
          + [N]
      )
      dist = model(vals, lengths)
      edges, enum_lengths = enumerate_support(dist)
      log_probs = dist.log_prob(edges)
      for b in range(lengths.shape[0]):
          log_probs[enum_lengths[b] :, b] = -1e9
      assert torch.isclose(
          log_probs.exp().sum(0), torch.tensor(1.0)
      ).all()
      entropy = dist.entropy
      assert torch.isclose(
          entropy, -log_probs.exp().mul(log_probs).sum(0)
      ).all()
  
      vals2 = torch.rand(*vals.shape)
      dist2 = model(vals2, lengths)
  
      cross_entropy = dist.cross_entropy(other=dist2)
      kl = dist.kl(other=dist2)
  
      edges2, enum_lengths2 = enumerate_support(dist2)
      log_probs2 = dist2.log_prob(edges2)
      for b in range(lengths.shape[0]):
          log_probs2[enum_lengths2[b] :, b] = -1e9
  
      assert torch.isclose(
          cross_entropy, -log_probs.exp().mul(log_probs2).sum(0)
      ).all()
      assert torch.isclose(
          kl, -log_probs.exp().mul(log_probs2 - log_probs).sum(0)
      ).all()
  
      argmax = dist.argmax
      _, max_indices = log_probs.max(0)
  
      amax = edges[max_indices, torch.arange(batch)]
      print(argmax.nonzero())
      print((amax - argmax).nonzero(), lengths)
      assert (amax == argmax).all()
  
      samples = dist.sample((100,))
      marginals = dist.marginals
      assert ((samples.mean(0) - marginals).abs() < 0.2).all()
  
      dist.kmax(5)
      dist.count
  
file_path: pytorch-struct/tests/test_distributions.py
imports: from torch_struct import LinearChainCRF, Autoregressive, KMaxSemiring, LogSemiring
  import torch
  from hypothesis import given, settings
  from hypothesis.strategies import integers, data, sampled_from
  from .extensions import test_lookup
label_desc: Use image node_id 2 for calling a global function with example usage: test_simple(data, seed) without return types