docstr: Test the example in the docstring of reorder_data_kjt
code: def test_reorder_data_kjt(self) -> None:
      """Test the example in the docstring of reorder_data_kjt"""
      keys = ["Key0", "Key1"]
      values = torch.arange(7).float()
      weights = values / 10.0
      lengths = torch.tensor([2, 0, 1, 1, 1, 2])
  
      # With weights
      x = KeyedJaggedTensor(
          keys=keys, values=values, lengths=lengths, weights=weights
      )
      y = reorder_data_kjt(x, torch.tensor([2, 1, 0]))
      self.assertEqual(y.keys(), keys)
      assert torch.allclose(
          y.values(), torch.tensor([2.0, 0.0, 1.0, 5.0, 6.0, 4.0, 3.0])
      )
      assert torch.allclose(
          y.lengths(), torch.tensor([1, 0, 2, 2, 1, 1])
      )
      assert torch.allclose(y.weights(), y.values() / 10.0)
  
      # Without weights
      x = KeyedJaggedTensor(keys=keys, values=values, lengths=lengths)
      y = reorder_data_kjt(x, torch.tensor([2, 1, 0]))
      self.assertEqual(y.keys(), keys)
      assert torch.allclose(
          y.values(), torch.tensor([2.0, 0.0, 1.0, 5.0, 6.0, 4.0, 3.0])
      )
      assert torch.allclose(
          y.lengths(), torch.tensor([1, 0, 2, 2, 1, 1])
      )
  
file_path: ReAgent/reagent/test/base/test_utils.py
imports: import unittest
  import numpy.testing
  import torch
  from reagent.core.torch_utils import masked_softmax, reorder_data_kjt, rescale_torch_tensor, shift_kjt_by_one, split_sequence_keyed_jagged_tensor
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 4 for calling the TestUtils obj's underlying member method code with example usage: obj.test_reorder_data_kjt() without return types