spikingjelly.activation_based.tensor_cache package

class spikingjelly.activation_based.cuda_kernel.tensor_cache.DataTypeConvertCUDACode

基类：object

float2bool = '\n    extern "C" __global__\n            void float2bool(const float* fs, unsigned char* bs, const int &N)\n            {\n                // assert N == numel / 8 and numel % 8 == 0\n                const int index = blockIdx.x * blockDim.x + threadIdx.x;\n                if (index < N)\n                {\n                    bs[index] = 0;\n                    const int mem_offset = (index << 3);\n                    #pragma unroll\n                    for(int i = 0; i < 8; i++)\n                    {\n                        bs[index] += ( ((unsigned char) fs[mem_offset + i]) << i);\n                    }\n                }\n            }\n    '

half2bool = '\n    #include <cuda_fp16.h>\n    extern "C" __global__\n            void half2bool(const half* fs, unsigned char* bs, const int &N)\n            {\n                // assert N == numel / 8 and numel % 8 == 0\n                const int index = blockIdx.x * blockDim.x + threadIdx.x;\n                if (index < N)\n                {\n                    bs[index] = 0;\n                    const int mem_offset = (index << 3);\n                    #pragma unroll\n                    for(int i = 0; i < 8; i++)\n                    {\n                        bs[index] += ( ((unsigned char) __half2float(fs[mem_offset + i])) << i);\n                    }\n                }\n            }\n    '

bool2float = '\n    extern "C" __global__\n            void bool2float(const unsigned char* bs, float* fs, const int &N)\n            {\n                const int index = blockIdx.x * blockDim.x + threadIdx.x;\n                if (index < N)\n                {\n                    const int mem_offset = (index << 3);\n                    unsigned char compressed_v = bs[index];\n                    #pragma unroll\n                    for(int i = 0; i < 8; i++)\n                    {\n                        fs[mem_offset + i] = (float) (compressed_v % 2);\n                        compressed_v = (compressed_v >> 1);\n                    }\n                }\n            }\n    '

bool2half = '\n    #include <cuda_fp16.h>\n    extern "C" __global__\n            void bool2half(const unsigned char* bs, half* fs, const int &N)\n            {\n                const int index = blockIdx.x * blockDim.x + threadIdx.x;\n                if (index < N)\n                {\n                    const int mem_offset = (index << 3);\n                    unsigned char compressed_v = bs[index];\n                    #pragma unroll\n                    for(int i = 0; i < 8; i++)\n                    {\n                        fs[mem_offset + i] = __float2half((float) (compressed_v % 2));\n                        compressed_v = (compressed_v >> 1);\n                    }\n                }\n            }\n    '

spikingjelly.activation_based.cuda_kernel.tensor_cache.float_spike_to_bool(spike: Tensor)

API Language: 中文 | English

---

• 中文

将浮点脉冲张量 spike 压缩为 torch.uint8 张量。压缩后的每个元素保存原始 脉冲中的8个元素（二值位打包）。

参数:

spike (torch.Tensor) -- 浮点脉冲张量， dtype 必须为 torch.float 或 torch.half ， 且元素取值为0或1

返回:

(spike_b, s_dtype, s_shape, s_padding)

• spike_b: 压缩后的张量， dtype=torch.uint8 ，每个元素保存8个脉冲

• s_dtype: 原始脉冲张量的数据类型

• s_shape: 原始脉冲张量的形状

• s_padding: 为对齐到8的倍数所填充的元素数量

返回类型:

tuple

---

• English

Compress a floating-point spike tensor spike into a torch.uint8 tensor. Each element in the compressed tensor stores 8 binary spikes (bit packing).

参数:

spike (torch.Tensor) -- Spike tensor whose dtype is torch.float or torch.half, and all elements are 0 or 1

返回:

(spike_b, s_dtype, s_shape, s_padding)

• spike_b: Compressed spike tensor with dtype=torch.uint8 and each element storing 8 spikes

• s_dtype: Dtype of the original spike tensor

• s_shape: Shape of the original spike tensor

• s_padding: Number of padded elements used for 8-element alignment

返回类型:

tuple

spikingjelly.activation_based.cuda_kernel.tensor_cache.bool_spike_to_float(spike_b: Tensor, s_dtype: dtype, s_shape: Size, s_padding: int = 0)

API Language: 中文 | English

---

• 中文

将压缩后的 torch.uint8 脉冲张量解压为原始浮点脉冲张量， 并按 s_shape 恢复形状。

参数:

• spike_b (torch.Tensor) -- 压缩后的脉冲张量， dtype=torch.uint8 ， 每个元素保存8个脉冲

• s_dtype (torch.dtype) -- 原始脉冲张量的数据类型

• s_shape (torch.Size) -- 原始脉冲张量的形状

• s_padding (int) -- 压缩时为对齐到8的倍数而填充的元素数量

返回:

解压并恢复形状后的原始脉冲张量

返回类型:

torch.Tensor

---

• English

Decompress a packed torch.uint8 spike tensor back to the original floating-point spike tensor and reshape it with s_shape.

参数:

• spike_b (torch.Tensor) -- Compressed spike tensor with dtype=torch.uint8 and each element storing 8 spikes

• s_dtype (torch.dtype) -- Dtype of the original spike tensor

• s_shape (torch.Size) -- Shape of the original spike tensor

• s_padding (int) -- Number of padded elements used for 8-element alignment during compression

返回:

The decompressed spike tensor restored to s_shape

返回类型:

torch.Tensor

spikingjelly.activation_based.cuda_kernel.tensor_cache.tensor_key(x: Tensor)

class spikingjelly.activation_based.cuda_kernel.tensor_cache.BoolTensorCache

基类：object

store_bool(spike: FloatTensor | HalfTensor)

get_float(tk, spike_shape: Size)