def pack(self, x) -> _Tuple[str, _Callable, _torch.device, _torch.Tensor]:
        """
        Function that will be called every time an operation saves a tensor for backward.
        """
        key = None
        op = lambda z: z.view(size)
        if x.numel() <= self.min_size:
            return x

        x_clone = x.clone() if self.max_mem <= 0 else None
        device = x.device
        size = x.size()

        if x.dtype.is_floating_point:
            grad_fn_list = []
            full_grad_fn_list = []
            c_grad_fn = x.grad_fn

            while len(grad_fn_list) < 2:
                if c_grad_fn:
                    str_grad_fn = str(type(c_grad_fn))

                    full_grad_fn_list.append(str_grad_fn)

                    if (
                        self.__class__.RESHAPE_BACKWARD in str_grad_fn
                        or self.__class__.TOCOPY_BACKWARD in str_grad_fn
                        or self.__class__.EXPAND_BACKWARD in str_grad_fn
                    ):
                        pass
                    else:
                        grad_fn_list.append(str_grad_fn)

                    c_grad_fn = c_grad_fn.next_functions[0][0] if c_grad_fn.next_functions else None
                else:
                    break

            if key is None:
                for _ in range(len(grad_fn_list), 2):
                    grad_fn_list.append("None")

                if (
                    self.__class__.SOFTMAX_BACKWARD in grad_fn_list[0]
                    and self.__class__.DIV_BACKWARD in grad_fn_list[1]
                ):
                    key = "softmax" + f".{self.pack_counter}"
                elif (
                    self.__class__.CLAMP_BACKWARD in grad_fn_list[0]
                    and self.__class__.SOFTMAX_BACKWARD in grad_fn_list[1]
                ):
                    key = "softmax" + f".{self.pack_counter}"
                    op = lambda z: z.view(size).clamp(min=self.zero_threshold)
                elif self.__class__.DIST_BACKWARD in grad_fn_list[0]:
                    self.pack_counter += 1
                    key = "x_c_dist" + f".{self.pack_counter}"
                elif (
                    self.__class__.VIEW_BACKWARD in grad_fn_list[0]
                    and self.__class__.DIST_BACKWARD in grad_fn_list[1]
                ):
                    key = "x_c_dist" + f".{self.pack_counter}"
                elif (
                    (
                        self.__class__.VIEW_BACKWARD in grad_fn_list[0]
                        and self.__class__.STACK_BACKWARD in grad_fn_list[1]
                    )
                    or (
                        self.__class__.STACK_BACKWARD in grad_fn_list[0]
                        and self.__class__.INDEX_BACKWARD in grad_fn_list[1]
                    )
                    or (
                        self.__class__.STACK_BACKWARD in grad_fn_list[0]
                        and self.__class__.SLICE_BACKWARD in grad_fn_list[1]
                    )
                ):
                    key = "X.b" + f".{-1}"
                elif (
                    self.__class__.TRANS_BACKWARD in grad_fn_list[0]
                    and self.__class__.STACK_BACKWARD in grad_fn_list[1]
                ):
                    key = "X.b" + f".{-1}"
                    if key in self.tensor_dict:
                        size = x.mT.size()
                        op = lambda z: z.reshape(size).mT
                    else:
                        key = None

            if key is None:
                key = self._compress_tensor(x, x.dtype)
            elif key not in self.tensor_dict:
                w = self._compress_tensor(x, x.dtype)
                self.tensor_dict[key] = w
        else:
            key = self._compress_tensor(x, _torch.uint8)
            op = lambda z: z.to(device, _torch.int32)

        return key, op, device, x_clone

        Function that will be called every time an operation saves a tensor for backward.
        
coremltools/optimize/torch/palettization/_fake_palettizer_tensor_hook.py