code: def test_batch_bucketized_one_hot(self, x, seed, gc, dc):
      np.random.seed(seed)
      d = x.shape[1]
      lens = np.random.randint(low=1, high=5, size=d)
      boundaries = []
      for i in range(d):
          # add [0, 0] as duplicated boundary for duplicated bucketization
          if lens[i] > 2:
              cur_boundary = np.append(
                  np.random.randn(lens[i] - 2) * 5, [0, 0]
              )
          else:
              cur_boundary = np.random.randn(lens[i]) * 5
          cur_boundary.sort()
          boundaries += cur_boundary.tolist()
  
      lens = np.array(lens, dtype=np.int32)
      boundaries = np.array(boundaries, dtype=np.float32)
  
      def ref(x, lens, boundaries):
          output_dim = lens.size + boundaries.size
          ret = np.zeros((x.shape[0], output_dim)).astype(x.dtype)
          boundary_offset = 0
          output_offset = 0
          for i, l in enumerate(lens):
              bucket_idx_right = np.digitize(
                  x[:, i],
                  boundaries[boundary_offset : boundary_offset + l],
                  right=True,
              )
              bucket_idx_left = np.digitize(
                  x[:, i],
                  boundaries[boundary_offset : boundary_offset + l],
                  right=False,
              )
              bucket_idx = np.floor_divide(
                  np.add(bucket_idx_right, bucket_idx_left), 2
              )
              for j in range(x.shape[0]):
                  ret[j, output_offset + bucket_idx[j]] = 1.0
              boundary_offset += lens[i]
              output_offset += lens[i] + 1
          return (ret,)
  
      op = core.CreateOperator(
          "BatchBucketOneHot", ["X", "LENS", "BOUNDARIES"], ["Y"]
      )
      self.assertReferenceChecks(gc, op, [x, lens, boundaries], 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 2 for calling the TestOneHotOps obj's underlying member method code with example usage: obj.test_batch_bucketized_one_hot(x, seed, gc, dc) and returns: ret