code: def test_segment_one_hot(self, hot_indices):
      index_size, lengths, indices = hot_indices
  
      index_size = np.array(index_size, dtype=np.int64)
      lengths = np.array(lengths, dtype=np.int32)
      indices = np.array(indices, dtype=np.int64)
  
      def segment_one_hot_ref(lengths, hot_indices, size):
          offset = 0
          out = np.zeros([len(lengths), size], dtype=float)
          for i, length in enumerate(lengths):
              for idx in hot_indices[offset : offset + length]:
                  out[i, idx] = 1.0
              offset += length
          return (out,)
  
      op = core.CreateOperator(
          "SegmentOneHot",
          ["lengths", "hot_indices", "size"],
          ["output"],
      )
      self.assertReferenceChecks(
          hu.cpu_do,
          op,
          [lengths, indices, index_size],
          segment_one_hot_ref,
      )
  
file_path: pytorch/caffe2/python/operator_test/one_hot_ops_test.py
imports: from caffe2.python import core, workspace
  from caffe2.proto import caffe2_pb2
  from hypothesis import given, settings
  import caffe2.python.hypothesis_test_util
  import caffe2.python.serialized_test.serialized_test_util
  import hypothesis.strategies
  import numpy
label_desc: Use image node_id 4 for calling the TestOneHotOps obj's underlying member method code with example usage: obj.test_segment_one_hot(hot_indices) and returns: out