import tensorflow as tf


class Compressor(object):
    """Interface for compressing and decompressing a given tensor."""
    @staticmethod
    def compress(tensor):
        """Compresses a tensor and returns it with the context needed to decompress it."""
        pass

    @staticmethod
    def decompress(tensor, ctx):
        """Decompress the tensor with the given context."""
        pass


class NoneCompressor(Compressor):
    """Default no-op compression."""
    @staticmethod
    def compress(tensor):
        """Returns the tensor unmodified."""
        return tensor, None

    @staticmethod
    def decompress(tensor, ctx):
        """Returns the tensor unmodified."""
        return tensor


class Int32Compressor(Compressor):
    """Compress all integer gradients to 32-bit."""
    @staticmethod
    def compress(tensor):
        """Downcasts the tensor to 32-bit."""
        tensor_compressed = tensor
        if tensor.dtype.is_integer:
            tensor_compressed = tf.cast(tensor, dtype=tf.int32)
        return tensor_compressed, tensor.dtype

    @staticmethod
    def decompress(tensor, dtype):
        """Upcasts the tensor to the initialization dtype."""
        tensor_decompressed = tensor
        if dtype.is_integer:
            tensor_decompressed = tf.cast(tensor, dtype=dtype)
        return tensor_decompressed


class FP16Compressor(Compressor):
    """Compress all floating point gradients to 16-bit."""
    @staticmethod
    def compress(tensor):
        """Downcasts the tensor to 16-bit."""
        tensor_compressed = tensor
        if tensor.dtype.is_floating:
            # Only allow compression from other floating point types
            tensor_compressed = tf.cast(tensor, dtype=tf.float16)
        return tensor_compressed, tensor.dtype

    @staticmethod
    def decompress(tensor, dtype):
        """Upcasts the tensor to the initialization dtype."""
        tensor_decompressed = tensor
        if dtype.is_floating:
            tensor_decompressed = tf.cast(tensor, dtype=dtype)
        return tensor_decompressed


class Compression(object):
    """Optional gradient compression algorithm used during allreduce."""

    """Do not compress the gradients. This is the default."""
    none = NoneCompressor

    """Compress all floating point gradients to 16-bit."""
    fp16 = FP16Compressor

    """Compress all integer gradients to 32-bits."""
    int32 = Int32Compressor

import tensorflow as tf


class Compressor(object):
    """Interface for compressing and decompressing a given tensor."""
    @staticmethod
    def compress(tensor):
        """Compresses a tensor and returns it with the context needed to decompress it."""
        pass

    @staticmethod
    def decompress(tensor, ctx):
        """Decompress the tensor with the given context."""
        pass


class NoneCompressor(Compressor):
    """Default no-op compression."""
    @staticmethod
    def compress(tensor):
        """Returns the tensor unmodified."""
        return tensor, None

    @staticmethod
    def decompress(tensor, ctx):
        """Returns the tensor unmodified."""
        return tensor


class Int32Compressor(Compressor):
    """Compress all integer gradients to 32-bit."""
    @staticmethod
    def compress(tensor):
        """Downcasts the tensor to 32-bit."""
        tensor_compressed = tensor
        if tensor.dtype.is_integer:
            tensor_compressed = tf.cast(tensor, dtype=tf.int32)
        return tensor_compressed, tensor.dtype

    @staticmethod
    def decompress(tensor, dtype):
        """Upcasts the tensor to the initialization dtype."""
        tensor_decompressed = tensor
        if dtype.is_integer:
            tensor_decompressed = tf.cast(tensor, dtype=dtype)
        return tensor_decompressed


class FP16Compressor(Compressor):
    """Compress all floating point gradients to 16-bit."""
    @staticmethod
    def compress(tensor):
        """Downcasts the tensor to 16-bit."""
        tensor_compressed = tensor
        if tensor.dtype.is_floating:
            # Only allow compression from other floating point types
            tensor_compressed = tf.cast(tensor, dtype=tf.float16)
        return tensor_compressed, tensor.dtype

    @staticmethod
    def decompress(tensor, dtype):
        """Upcasts the tensor to the initialization dtype."""
        tensor_decompressed = tensor
        if dtype.is_floating:
            tensor_decompressed = tf.cast(tensor, dtype=dtype)
        return tensor_decompressed


class Compression(object):
    """Optional gradient compression algorithm used during allreduce."""

    """Do not compress the gradients. This is the default."""
    none = NoneCompressor

    """Compress all floating point gradients to 16-bit."""
    fp16 = FP16Compressor

    """Compress all integer gradients to 32-bits."""
    int32 = Int32Compressor

Interface for compressing and decompressing a given tensor. Compresses a tensor and returns it with the context needed to decompress it. Decompress the tensor with the given context. Default no-op compression. Returns the tensor unmodified. Returns the tensor unmodified. Compress all integer gradients to 32-bit. Downcasts the tensor to 32-bit. Upcasts the tensor to the initialization dtype. Compress all floating point gradients to 16-bit. Downcasts the tensor to 16-bit. # Only allow compression from other floating point types Upcasts the tensor to the initialization dtype. Optional gradient compression algorithm used during allreduce. Do not compress the gradients. This is the default. Compress all floating point gradients to 16-bit. Compress all integer gradients to 32-bits.