code: def make_kjt(
      values: List[int], lengths: List[int]
  ) -> KeyedJaggedTensor:
      values_tensor = torch.tensor(values, dtype=torch.int32)
      lengths_tensor = torch.tensor(lengths, dtype=torch.int32)
      torch._check(
          torch.sum(lengths_tensor).item() == values_tensor.size(0)
      )
      kjt = KeyedJaggedTensor(
          keys=[f"key{i}" for i in range(len(lengths))],
          values=values_tensor,
          lengths=lengths_tensor,
      )
      dynamic_dim(kjt._values, 0)
      return kjt
  
imports: import sys
  import unittest
  from typing import List, Tuple
  import torch
  import torch._dynamo.skipfiles
  from fbgemm_gpu import sparse_ops
  from torch._export import dynamic_dim
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.test_utils.infer_utils import assert_close, create_test_model_ebc_only, dynamo_skipfiles_allow, KJTInputExportWrapper, prep_inputs, replace_registered_tbes_with_mock_tbes, replace_sharded_quant_modules_tbes_with_mock_tbes, TestQuantEBCSharder
  from torchrec.distributed.types import ShardingEnv, ShardingType
  from torchrec.sparse.jagged_tensor import ComputeKJTToJTDict, KeyedJaggedTensor
label_desc: Use image node_id 1 for calling a global function with example usage: make_kjt(values, lengths) and returns: kjt