jactorch.functional.quantization

Basic quantization functions with a straight-through gradient estimator.

Functions

quantize(x)	Quantize a tensor to binary values: (x > 0.5).float().
randomized_quantize(x)	Quantize a tensor to binary values: (rand() > x).float().

Functions

quantize(x)

Quantize a tensor to binary values: (x > 0.5).float(). This function implements the straight-through gradient estimator.

Parameters:

x (Tensor) – the input tensor.

Returns:

the quantized tensor.

Return type:

Tensor

randomized_quantize(x)

Quantize a tensor to binary values: (rand() > x).float(). This function implements the straight-through gradient estimator.

Parameters:

x (Tensor) – the input tensor.

Returns:

the quantized tensor.

Return type:

Tensor