text: import torch
  import numpy as np
  from torch.nn import Parameter
  from halp.layers.embedding import BitCenterEmbedding, bit_center_embedding
  from halp.utils.utils import void_cast_func, single_to_half_det, single_to_half_stoc
  from unittest import TestCase
  from halp.utils.test_utils import HalpTest
  from torch.autograd.gradcheck import get_numerical_jacobian, iter_tensors, make_jacobian
  from halp.layers.bit_center_layer_test import TestBitCenterLayer
  import logging
  import sys
  logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
  logger = logging.getLogger()
  
  
  class TestBitCenterEmbedding(TestBitCenterLayer, TestCase):
      '''
      Test the functionality of bit centering linear layers
      '''
      def get_config(self, type="grad_check"):
          config = {}
          if type == "grad_check":
              config["n_train_sample"] = 35
              config["num_embeddings"] = 17
              config["embedding_dim"] = 24
              config["bias"] = True
              config["cast_func"] = void_cast_func
              config["do_double"] = True
              config["seed"] = 0
              config["batch_size"] = 35
              config["seq_length"] = 11
          elif type == "fw_bw_proc":
              config["n_train_sample"] = 98
              config["num_embeddings"] = 13
              config["embedding_dim"] = 31
              config["bias"] = False
              config["cast_func"] = single_to_half_det
              config["do_double"] = False
              config["seed"] = 0
              config["batch_size"] = 33
              config["seq_length"] = 11
          else:
              raise Exception("Config type not supported!")
          return config
  
      def prepare_layer(self,
                        n_train_sample,
                        num_embeddings,
                        embedding_dim,
                        bias=False, # bias is dummy to adapt to the test interface
                        cast_func=void_cast_func,
                        do_double=True,
                        seed=0,
                        batch_size=1,
                        seq_length=10):
          layer = BitCenterEmbedding(
              num_embeddings=num_embeddings,
              embedding_dim=embedding_dim,
              cast_func=cast_func,
              n_train_sample=n_train_sample)
          # Note do_double = setup layer for gradient check, otherwise, it is for checking
          # the tensor properties
          if do_double:
              layer.double()
              layer.weight.data.copy_(
                  torch.randn(
                      num_embeddings, embedding_dim, dtype=torch.double,
                      requires_grad=False).cuda())
              layer.weight_lp.data.copy_(layer.weight.data)
              layer.weight_delta.data.copy_(
                  torch.randn(
                      num_embeddings, embedding_dim, dtype=torch.double,
                      requires_grad=True).cuda())
          layer.cuda()
          return layer
  
      def get_input(self,
                    n_train_sample,
                    num_embeddings,
                    embedding_dim,
                    bias=False,
                    cast_func=void_cast_func,
                    do_double=True,
                    seed=0,
                    batch_size=1,
                    seq_length=10):
          np.random.seed(seed)
          torch.manual_seed(seed)
          torch.cuda.manual_seed_all(seed)
          self.target_dtype = None
  
          input_delta = Parameter(
              torch.zeros(n_train_sample, seq_length, dtype=torch.long).cuda().random_(0, num_embeddings),
              requires_grad=False)
          input_fp = Parameter(
              torch.zeros(n_train_sample, seq_length, dtype=torch.long).cuda().data.copy_(input_delta.data),
              requires_grad=False)
          return [
              input_fp,
          ], [
              input_delta,
          ]
  
      def check_layer_param_and_cache(self, layer):
          t_list = [(layer.weight, torch.float32, True, True),
                    (layer.weight_delta, torch.half, True, True),
                    (layer.weight_lp, torch.half, True, False),
                    (layer.input_cache, torch.long, False, False),
                    (layer.grad_output_cache, torch.half, False, False)]
          self.CheckLayerTensorProperty(t_list)
          self.CheckLayerTensorGradProperty(t_list)
  
      def get_numerical_grad(self,
                             layer,
                             input_fp,
                             input_delta,
                             perturb_eps,
                             target=None):
          grad_list = []
          layer.set_mode(do_offset=True)
          param_dict = layer.state_dict()
          # update the offset variable
          for name, param in layer.named_parameters():
              if name.endswith("_delta"):
                  p_offset = param_dict[name.split("_delta")[0]]
                  p_offset.data.add_(param)
          
          input_final = [  # for embedding the input_fp and input_delta are the same index array
              Parameter(x, requires_grad=False)
              for x, y in zip(input_fp, input_delta)
          ]
          output_final = layer(*input_final)
          loss = 0.5 * torch.sum(output_final**2)
          loss.backward()
          # use the gradient from 0.5*sum(output**2), this case has output = gradient wrt output
          num_input_grad = None
          grad_list.append(num_input_grad)
  
          def get_loss(x):
              output = layer(*x)
              return torch.sum(0.5 * output * output)
  
          grad_list += self.get_numerical_param_grad(layer, input_final, get_loss,
                                                     perturb_eps)
          return output_final, grad_list
  
  
  if __name__ == "__main__":
      print(torch.__version__)
      unittest.main()