def test_bad_input_bits_raises(self):
with self.assertRaisesRegexp(TypeError, 'cannot be a TensorFlow value'):
stages_impl.BitpackingEncodingStage(tf.constant(1, dtype=tf.int32))
with self.assertRaisesRegexp(ValueError, 'between 1 and 16'):
stages_impl.BitpackingEncodingStage(0)
with self.assertRaisesRegexp(ValueError, 'between 1 and 16'):
stages_impl.BitpackingEncodingStage(17)
def hadamard_quantization(bits):
"""Returns hadamard quanitzation `Encoder`.
The `Encoder` first reshapes the input to a rank-1 `Tensor`, and applies the
Hadamard transform (rotation). It then applies uniform quantization with the
extreme values being the minimum and maximum of the rotated vector being
encoded. Finally, the quantized values are bitpacked to an integer type.
The `Encoder` is a composition of the following encoding stages:
* `FlattenEncodingStage` - reshaping the input to a vector.
* `HadamardEncodingStage` - applying the Hadamard transform.
* `UniformQuantizationEncodingStage` - applying uniform quantization.
* `BitpackingEncodingStage` - bitpacking the result into integer values.
Args:
bits: Number of bits to quantize into.
Returns:
The hadamard quantization `Encoder`.
"""
return core_encoder.EncoderComposer(
stages_impl.BitpackingEncodingStage(bits)).add_parent(
stages_impl.UniformQuantizationEncodingStage(bits),
stages_impl.UniformQuantizationEncodingStage.ENCODED_VALUES_KEY
).add_parent(
stages_impl.HadamardEncodingStage(),
stages_impl.HadamardEncodingStage.ENCODED_VALUES_KEY).add_parent(
stages_impl.FlattenEncodingStage(),
stages_impl.FlattenEncodingStage.ENCODED_VALUES_KEY).make()
def uniform_quantization(bits):
"""Returns uniform quanitzation `Encoder`.
The `Encoder` first reshapes the input to a rank-1 `Tensor`, then applies
uniform quantization with the extreme values being the minimum and maximum of
the vector being encoded. Finally, the quantized values are bitpacked to an
integer type.
The `Encoder` is a composition of the following encoding stages:
* `FlattenEncodingStage`
* `UniformQuantizationEncodingStage`
* `BitpackingEncodingStage`
Args:
bits: Number of bits to quantize into.
Returns:
The quantization `Encoder`.
"""
return core_encoder.EncoderComposer(
stages_impl.BitpackingEncodingStage(bits)).add_parent(
stages_impl.UniformQuantizationEncodingStage(bits), stages_impl.
UniformQuantizationEncodingStage.ENCODED_VALUES_KEY).add_parent(
stages_impl.FlattenEncodingStage(),
stages_impl.FlattenEncodingStage.ENCODED_VALUES_KEY).make()
def test_bad_input_executes(self):
# Test that if input to encode is outside of the expected range, everything
# still executes, but the result is not correct.
x = np.array([2**9] * 5).astype(np.int32)
stage = stages_impl.BitpackingEncodingStage(8)
test_data = self.run_one_to_many_encode_decode(
stage, lambda: tf.constant(x, tf.float32))
self.assertNotAllClose(x, test_data.decoded_x.astype(np.int32))
def test_encoded_values_as_expected(self, bits, expected_bitpacked_values):
# Tests that the packed values are as expected.
stage = stages_impl.BitpackingEncodingStage(bits)
test_data = self.run_one_to_many_encode_decode(
stage, lambda: tf.constant([1.0, 0.0, 1.0, 1.0, 0.0], dtype=tf.float32))
self.assertAllEqual(
expected_bitpacked_values, test_data.encoded_x[
stages_impl.BitpackingEncodingStage.ENCODED_VALUES_KEY])
def test_is_lossless(self, bits, shape):
# Tests that the encoding is lossless, for a variety of inputs.
def x_fn():
return tf.cast(
tf.random.uniform(
shape, minval=0, maxval=2**bits - 1, dtype=tf.int32), tf.float32)
stage = stages_impl.BitpackingEncodingStage(bits)
test_data = self.run_one_to_many_encode_decode(stage, x_fn)
self.assertAllClose(test_data.x, test_data.decoded_x)
self.assertEqual(test_data.x.dtype, test_data.decoded_x.dtype)
def default_encoding_stage(self): """See base class.""" return stages_impl.BitpackingEncodingStage(8)