torch-floating-point 0.0.1__tar.gz

torch_floating_point-0.0.1/PKG-INFO

@@ -0,0 +1,274 @@
+Metadata-Version: 2.4
+Name: torch-floating-point
+Version: 0.0.1
+Summary: A PyTorch library for custom floating point quantization with autograd support
+Author-email: Samir Moustafa <samir.moustafa.97@gmail.com>
+Maintainer-email: Samir Moustafa <samir.moustafa.97@gmail.com>
+Project-URL: Homepage, https://github.com/SamirMoustafa/torch-floating-point
+Project-URL: Documentation, https://torch-floating-point.readthedocs.io/
+Project-URL: Repository, https://github.com/SamirMoustafa/torch-floating-point
+Project-URL: Bug Tracker, https://github.com/SamirMoustafa/torch-floating-point/issues
+Project-URL: Source Code, https://github.com/SamirMoustafa/torch-floating-point
+Keywords: pytorch,floating-point,quantization,autograd,machine-learning,deep-learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Framework :: Pytest
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Provides-Extra: dev
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: pre-commit>=3.0.0; extra == "dev"
+
+# Torch Floating Point
+
+A PyTorch library for custom floating point quantization with autograd support. This library provides efficient implementations of custom floating point formats with automatic differentiation capabilities.
+
+## Features
+
+- **Custom Floating Point Formats**: Support for arbitrary floating point configurations (sign bits, exponent bits, mantissa bits, bias)
+- **Autograd Support**: Full PyTorch autograd integration for training with quantized weights
+- **CUDA Support**: GPU acceleration for both forward and backward passes
+- **Multiple Precision**: Support for various bit widths (4-bit, 8-bit, 16-bit, 32-bit)
+- **Straight-Through Estimator**: Gradient-friendly quantization for training
+- **Comprehensive Testing**: Extensive test suite covering differentiability and accuracy
+
+## Installation
+
+### From PyPI (Recommended)
+
+```bash
+pip install torch-floating-point
+```
+
+### From Source
+
+```bash
+git clone https://github.com/SamirMoustafa/torch-floating-point.git
+cd torch-floating-point
+pip install -e .
+```
+
+### Development Installation
+
+```bash
+git clone https://github.com/SamirMoustafa/torch-floating-point.git
+cd torch-floating-point
+pip install -e ".[dev,test]"
+pre-commit install
+```
+
+## Quick Start
+
+```python
+import torch
+from floating_point import FloatingPoint, Round
+
+# Define a custom 8-bit floating point format (1 sign, 4 exponent, 3 mantissa bits)
+fp8 = FloatingPoint(sign_bits=1, exponent_bits=4, mantissa_bits=3, bias=7, bits=8)
+
+# Create a rounding function
+rounder = Round(fp8)
+
+# Create a tensor with gradients
+x = torch.randn(10, requires_grad=True)
+
+# Quantize the tensor
+quantized = rounder(x)
+
+# Use in training (gradients flow through)
+loss = quantized.sum()
+loss.backward()
+
+print(f"Original: {x}")
+print(f"Quantized: {quantized}")
+print(f"Gradients: {x.grad}")
+```
+
+## Usage Examples
+
+### Custom Floating Point Configuration
+
+```python
+from floating_point import FloatingPoint
+
+# 4-bit floating point (1 sign, 2 exponent, 1 mantissa)
+fp4 = FloatingPoint(sign_bits=1, exponent_bits=2, mantissa_bits=1, bias=1, bits=4)
+
+# 8-bit floating point with custom max mantissa
+fp8_custom = FloatingPoint(
+    sign_bits=1, 
+    exponent_bits=4, 
+    mantissa_bits=3, 
+    bias=7, 
+    bits=8,
+    max_mantissa_at_max_exponent=6,  # Custom max mantissa
+    reserved_exponent=False  # No reserved exponent for inf/nan
+)
+
+# 16-bit floating point (standard)
+fp16 = FloatingPoint(sign_bits=1, exponent_bits=5, mantissa_bits=10, bias=15, bits=16)
+```
+
+### Training with Quantized Weights
+
+```python
+import torch
+import torch.nn as nn
+from floating_point import FloatingPoint, Round
+
+class QuantizedLinear(nn.Module):
+    def __init__(self, in_features, out_features, fp_config):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(out_features, in_features))
+        self.rounder = Round(fp_config)
+    
+    def forward(self, x):
+        quantized_weight = self.rounder(self.weight)
+        return torch.nn.functional.linear(x, quantized_weight)
+
+# Define quantization format
+fp8 = FloatingPoint(sign_bits=1, exponent_bits=4, mantissa_bits=3, bias=7, bits=8)
+
+# Create model with quantized weights
+model = QuantizedLinear(784, 10, fp8)
+optimizer = torch.optim.Adam(model.parameters())
+
+# Training loop
+for epoch in range(10):
+    # ... your training code ...
+    loss.backward()
+    optimizer.step()
+```
+
+### Direct Function Usage
+
+```python
+import torch
+from floating_point import autograd
+
+# Direct quantization function
+x = torch.randn(100, requires_grad=True)
+quantized = autograd(x, exponent_bits=4, mantissa_bits=3, bias=7)
+
+# Gradients work automatically
+loss = quantized.sum()
+loss.backward()
+```
+
+## Supported Formats
+
+The library supports various floating point formats:
+
+| Format | Sign Bits | Exponent Bits | Mantissa Bits | Bias | Total Bits |
+|--------|-----------|---------------|---------------|------|------------|
+| FP4    | 1         | 2             | 1             | 1    | 4          |
+| FP8    | 1         | 4             | 3             | 7    | 8          |
+| FP16   | 1         | 5             | 10            | 15   | 16         |
+| BF16   | 1         | 8             | 7             | 127  | 16         |
+| FP32   | 1         | 8             | 23            | 127  | 32         |
+
+## Development
+
+### Testing
+
+The project includes two testing approaches:
+
+1. **CI/CD Tests** (GitHub Actions): Fast, lightweight tests that verify core functionality without heavy numerical computations
+2. **Full Test Suite**: Complete test coverage including all numerical precision tests (run locally or via manual workflow trigger)
+
+To run the full test suite locally:
+```bash
+export LD_LIBRARY_PATH=$(python -c "import torch; print(torch.__file__)")/lib:$LD_LIBRARY_PATH
+python -m pytest test/round.py test/data_types.py -v
+```
+
+### Running Tests
+
+```bash
+# Run all tests
+make test
+
+# Run tests with coverage
+make test-cov
+
+# Run specific test file
+python -m pytest test/round.py -v
+```
+
+### Code Quality
+
+```bash
+# Run linting
+make lint
+
+# Format code
+make format
+
+# Run all checks
+make full-check
+```
+
+### Building
+
+```bash
+# Build extension
+cd floating_point && python setup.py build_ext --inplace
+
+# Build package
+make build
+
+# Clean build artifacts
+make clean
+```
+
+## Contributing
+
+1. Fork the repository
+2. Create a feature branch (`git checkout -b feature/amazing-feature`)
+3. Install development dependencies (`make setup-dev`)
+4. Make your changes
+5. Run tests (`make test`)
+6. Run linting (`make lint`)
+7. Commit your changes (`git commit -m 'Add amazing feature'`)
+8. Push to the branch (`git push origin feature/amazing-feature`)
+9. Open a Pull Request
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## Citation
+
+If you use this library in your research, please cite:
+
+```bibtex
+@software{torch_floating_point,
+  title={Torch Floating Point: A PyTorch library for custom floating point quantization},
+  author={Samir Moustafa},
+  year={2024},
+  url={https://github.com/SamirMoustafa/torch-floating-point}
+}
+```
+
+## Acknowledgments
+
+- PyTorch team for the excellent autograd system
+- The PyTorch C++ extension community for guidance on extension development
+- Contributors and users of this library
+
+## Support
+
+- **Issues**: [GitHub Issues](https://github.com/SamirMoustafa/torch-floating-point/issues)
+- **Discussions**: [GitHub Discussions](https://github.com/SamirMoustafa/torch-floating-point/discussions)
+- **Email**: samir.moustafa.97@gmail.com

torch_floating_point-0.0.1/README.md

@@ -0,0 +1,242 @@
+# Torch Floating Point
+
+A PyTorch library for custom floating point quantization with autograd support. This library provides efficient implementations of custom floating point formats with automatic differentiation capabilities.
+
+## Features
+
+- **Custom Floating Point Formats**: Support for arbitrary floating point configurations (sign bits, exponent bits, mantissa bits, bias)
+- **Autograd Support**: Full PyTorch autograd integration for training with quantized weights
+- **CUDA Support**: GPU acceleration for both forward and backward passes
+- **Multiple Precision**: Support for various bit widths (4-bit, 8-bit, 16-bit, 32-bit)
+- **Straight-Through Estimator**: Gradient-friendly quantization for training
+- **Comprehensive Testing**: Extensive test suite covering differentiability and accuracy
+
+## Installation
+
+### From PyPI (Recommended)
+
+```bash
+pip install torch-floating-point
+```
+
+### From Source
+
+```bash
+git clone https://github.com/SamirMoustafa/torch-floating-point.git
+cd torch-floating-point
+pip install -e .
+```
+
+### Development Installation
+
+```bash
+git clone https://github.com/SamirMoustafa/torch-floating-point.git
+cd torch-floating-point
+pip install -e ".[dev,test]"
+pre-commit install
+```
+
+## Quick Start
+
+```python
+import torch
+from floating_point import FloatingPoint, Round
+
+# Define a custom 8-bit floating point format (1 sign, 4 exponent, 3 mantissa bits)
+fp8 = FloatingPoint(sign_bits=1, exponent_bits=4, mantissa_bits=3, bias=7, bits=8)
+
+# Create a rounding function
+rounder = Round(fp8)
+
+# Create a tensor with gradients
+x = torch.randn(10, requires_grad=True)
+
+# Quantize the tensor
+quantized = rounder(x)
+
+# Use in training (gradients flow through)
+loss = quantized.sum()
+loss.backward()
+
+print(f"Original: {x}")
+print(f"Quantized: {quantized}")
+print(f"Gradients: {x.grad}")
+```
+
+## Usage Examples
+
+### Custom Floating Point Configuration
+
+```python
+from floating_point import FloatingPoint
+
+# 4-bit floating point (1 sign, 2 exponent, 1 mantissa)
+fp4 = FloatingPoint(sign_bits=1, exponent_bits=2, mantissa_bits=1, bias=1, bits=4)
+
+# 8-bit floating point with custom max mantissa
+fp8_custom = FloatingPoint(
+    sign_bits=1, 
+    exponent_bits=4, 
+    mantissa_bits=3, 
+    bias=7, 
+    bits=8,
+    max_mantissa_at_max_exponent=6,  # Custom max mantissa
+    reserved_exponent=False  # No reserved exponent for inf/nan
+)
+
+# 16-bit floating point (standard)
+fp16 = FloatingPoint(sign_bits=1, exponent_bits=5, mantissa_bits=10, bias=15, bits=16)
+```
+
+### Training with Quantized Weights
+
+```python
+import torch
+import torch.nn as nn
+from floating_point import FloatingPoint, Round
+
+class QuantizedLinear(nn.Module):
+    def __init__(self, in_features, out_features, fp_config):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(out_features, in_features))
+        self.rounder = Round(fp_config)
+    
+    def forward(self, x):
+        quantized_weight = self.rounder(self.weight)
+        return torch.nn.functional.linear(x, quantized_weight)
+
+# Define quantization format
+fp8 = FloatingPoint(sign_bits=1, exponent_bits=4, mantissa_bits=3, bias=7, bits=8)
+
+# Create model with quantized weights
+model = QuantizedLinear(784, 10, fp8)
+optimizer = torch.optim.Adam(model.parameters())
+
+# Training loop
+for epoch in range(10):
+    # ... your training code ...
+    loss.backward()
+    optimizer.step()
+```
+
+### Direct Function Usage
+
+```python
+import torch
+from floating_point import autograd
+
+# Direct quantization function
+x = torch.randn(100, requires_grad=True)
+quantized = autograd(x, exponent_bits=4, mantissa_bits=3, bias=7)
+
+# Gradients work automatically
+loss = quantized.sum()
+loss.backward()
+```
+
+## Supported Formats
+
+The library supports various floating point formats:
+
+| Format | Sign Bits | Exponent Bits | Mantissa Bits | Bias | Total Bits |
+|--------|-----------|---------------|---------------|------|------------|
+| FP4    | 1         | 2             | 1             | 1    | 4          |
+| FP8    | 1         | 4             | 3             | 7    | 8          |
+| FP16   | 1         | 5             | 10            | 15   | 16         |
+| BF16   | 1         | 8             | 7             | 127  | 16         |
+| FP32   | 1         | 8             | 23            | 127  | 32         |
+
+## Development
+
+### Testing
+
+The project includes two testing approaches:
+
+1. **CI/CD Tests** (GitHub Actions): Fast, lightweight tests that verify core functionality without heavy numerical computations
+2. **Full Test Suite**: Complete test coverage including all numerical precision tests (run locally or via manual workflow trigger)
+
+To run the full test suite locally:
+```bash
+export LD_LIBRARY_PATH=$(python -c "import torch; print(torch.__file__)")/lib:$LD_LIBRARY_PATH
+python -m pytest test/round.py test/data_types.py -v
+```
+
+### Running Tests
+
+```bash
+# Run all tests
+make test
+
+# Run tests with coverage
+make test-cov
+
+# Run specific test file
+python -m pytest test/round.py -v
+```
+
+### Code Quality
+
+```bash
+# Run linting
+make lint
+
+# Format code
+make format
+
+# Run all checks
+make full-check
+```
+
+### Building
+
+```bash
+# Build extension
+cd floating_point && python setup.py build_ext --inplace
+
+# Build package
+make build
+
+# Clean build artifacts
+make clean
+```
+
+## Contributing
+
+1. Fork the repository
+2. Create a feature branch (`git checkout -b feature/amazing-feature`)
+3. Install development dependencies (`make setup-dev`)
+4. Make your changes
+5. Run tests (`make test`)
+6. Run linting (`make lint`)
+7. Commit your changes (`git commit -m 'Add amazing feature'`)
+8. Push to the branch (`git push origin feature/amazing-feature`)
+9. Open a Pull Request
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## Citation
+
+If you use this library in your research, please cite:
+
+```bibtex
+@software{torch_floating_point,
+  title={Torch Floating Point: A PyTorch library for custom floating point quantization},
+  author={Samir Moustafa},
+  year={2024},
+  url={https://github.com/SamirMoustafa/torch-floating-point}
+}
+```
+
+## Acknowledgments
+
+- PyTorch team for the excellent autograd system
+- The PyTorch C++ extension community for guidance on extension development
+- Contributors and users of this library
+
+## Support
+
+- **Issues**: [GitHub Issues](https://github.com/SamirMoustafa/torch-floating-point/issues)
+- **Discussions**: [GitHub Discussions](https://github.com/SamirMoustafa/torch-floating-point/discussions)
+- **Email**: samir.moustafa.97@gmail.com

torch_floating_point-0.0.1/floating_point/__init__.py

@@ -0,0 +1 @@
+from .floating_point import autograd

torch_floating_point-0.0.1/floating_point/data_types.py

@@ -0,0 +1,140 @@
+import math
+from itertools import product
+from typing import List, Optional
+
+
+class FloatingPoint:
+    """
+    A class representing a custom floating-point format.
+    Parameters:
+    - sign_bits: Number of bits for the sign (0 or 1).
+    - exponent_bits: Number of bits for the exponent.
+    - mantissa_bits: Number of bits for the mantissa.
+    - bias: Bias for the exponent.
+    - bits: Total number of bits (sign + exponent + mantissa).
+    - max_mantissa_at_max_exponent: Maximum mantissa value at maximum exponent.
+    - reserved_exponent: Whether the exponent has reserved values (default is True).
+    """
+
+    def __init__(
+        self,
+        sign_bits: int,
+        exponent_bits: int,
+        mantissa_bits: int,
+        bias: int,
+        bits: int,
+        max_mantissa_at_max_exponent: Optional[int] = None,
+        reserved_exponent: bool = True,
+    ):
+        assert bits > 0, "Total bits must be positive."
+        assert bits == sign_bits + exponent_bits + mantissa_bits, (
+            "Sum of sign, exponent, and mantissa bits must equal bits."
+        )
+        assert 0 <= sign_bits <= 1, "Sign bits must be 0 or 1."
+        self.bits = bits
+        self.sign_bits = sign_bits
+        self.exponent_bits = exponent_bits
+        self.mantissa_bits = mantissa_bits
+        self.bias = bias
+        self.reserved_exponent = reserved_exponent
+        if max_mantissa_at_max_exponent is not None:
+            self.max_mantissa_at_max_exponent = max_mantissa_at_max_exponent
+        else:
+            self.max_mantissa_at_max_exponent = 2**mantissa_bits - 1
+
+    @property
+    def is_signed(self) -> bool:
+        return self.sign_bits > 0
+
+    @property
+    def epsilon(self) -> float:
+        # Smallest positive subnormal value
+        return float(2 ** (-self.mantissa_bits))
+
+    @property
+    def minimum(self) -> float:
+        # Negative of the max finite value
+        return -self.maximum if self.is_signed else 0.0
+
+    @property
+    def maximum(self) -> float:
+        if self.exponent_bits == 0:
+            # Subnormal maximum: (max_mantissa / 2^mantissa_bits) * 2^(1 - bias)
+            max_exponent = 1 - self.bias
+            max_mantissa = (2**self.mantissa_bits) - 1
+            return float((max_mantissa / (2**self.mantissa_bits)) * (2**max_exponent))
+        else:
+            # Calculate max stored exponent based on reserved status
+            max_stored_exponent = (2**self.exponent_bits - 2) if self.reserved_exponent else (2**self.exponent_bits - 1)
+            max_exponent = max_stored_exponent - self.bias
+            return float((1 + self.max_mantissa_at_max_exponent / (2**self.mantissa_bits)) * (2**max_exponent))
+
+    def generate_bit_combinations(self) -> List[int]:
+        """Generate all possible bit patterns for the given configuration."""
+        total_bits = self.bits
+        bit_combinations = list(product([0, 1], repeat=total_bits))
+        return [int("".join(map(str, bits)), 2) for bits in bit_combinations]
+
+    def bit_pattern_to_custom_fp(self, bit_pattern: int) -> float:
+        total_bits = self.sign_bits + self.exponent_bits + self.mantissa_bits
+        # Mask definitions
+        sign_mask = (1 << (total_bits - 1)) if self.is_signed else 0
+        exponent_mask = ((1 << self.exponent_bits) - 1) << self.mantissa_bits
+        mantissa_mask = (1 << self.mantissa_bits) - 1
+        sign = (bit_pattern & sign_mask) >> (self.exponent_bits + self.mantissa_bits) if self.is_signed else 0
+        exponent = (bit_pattern & exponent_mask) >> self.mantissa_bits
+        mantissa = bit_pattern & mantissa_mask
+        # Decode components
+        sign_factor = -1 if sign else 1
+        if self.exponent_bits == 0:
+            # Only subnormals
+            exponent_value = 1 - self.bias
+            mantissa_value = mantissa / (2**self.mantissa_bits)
+            if mantissa == 0:
+                return sign_factor * 0.0
+            return float(sign_factor * mantissa_value * (2**exponent_value))
+        else:
+            max_exponent = (1 << self.exponent_bits) - 1
+            if self.reserved_exponent and exponent == max_exponent:
+                if mantissa == 0:
+                    return sign_factor * math.inf
+                else:
+                    return math.nan
+            elif exponent == 0:
+                if mantissa == 0:
+                    return sign_factor * 0.0
+                else:
+                    mantissa_value = mantissa / (2**self.mantissa_bits)
+                    return float(sign_factor * mantissa_value * (2 ** (1 - self.bias)))
+            else:
+                exponent_value = exponent - self.bias
+                mantissa_value = 1 + (mantissa / (2**self.mantissa_bits))
+                return float(sign_factor * mantissa_value * (2**exponent_value))
+
+    def generate_all_custom_fp_values(self) -> List[float]:
+        bit_combinations = self.generate_bit_combinations()
+        values = [self.bit_pattern_to_custom_fp(b) for b in bit_combinations]
+        values = sorted(
+            values,
+            key=lambda x: (
+                math.inf if math.isnan(x) else math.copysign(1, x),
+                math.inf if math.isnan(x) else x,
+            ),
+        )
+        assert len(values) == 2**self.bits, f"Incorrect number of values generated: {len(values)} != {2**self.bits}"
+        return values
+
+    @property
+    def values(self) -> List[float]:
+        return self.generate_all_custom_fp_values()
+
+    def __repr__(self) -> str:
+        return (
+            f"Float{self.bits}-"
+            f"S{self.sign_bits}"
+            f"E{self.exponent_bits}"
+            f"M{self.mantissa_bits}"
+            f"B{self.bias}"
+            f"MaxM{self.max_mantissa_at_max_exponent}"
+            f"{'R' if self.reserved_exponent else ''}"
+        )