tf.tpu.experimental.embedding.QuantizationConfig

Settings for simulated quantization of the tpu embedding table.

tf.tpu.experimental.embedding.QuantizationConfig(
    num_buckets: int, lower: float, upper: float
)

When simulated quantization is enabled, the results of the embedding lookup are clipped and quantized according to the settings here before the combiner is applied.

For example, to quantize input the following is done:

if input < lower
  input = lower
if input > upper
  input = upper
quantum = (upper - lower) / (num_buckets - 1)
input = math.floor((input - lower) / quantum + 0.5) * quantium + lower

See tensorflow/core/protobuf/tpu/optimization_parameters.proto for more details.

num_buckets The number of quantization buckets, must be atleast 2.
lower The lower bound for the quantization range.
upper The upper bound for the quantization range.

ValueError if num_buckets is less than 2.