tessl/pypi-etils

Collection of common python utils for machine learning and scientific computing workflows

—

Pending

Overview

Eval results

Files

Array Type Annotations (array_types)

Name: tessl/pypi-etils
Author: tessl

Comprehensive type annotations for NumPy, JAX, TensorFlow, and PyTorch arrays with specific precision types for type-safe machine learning code development. Provides a unified typing system across different ML frameworks.

Capabilities

Core Array Types

Fundamental array type definitions for general use.

ArrayLike = Union[np.ndarray, list, tuple, int, float]
# Array-like objects including lists, tuples, scalars, and arrays

Array = ArrayAliasMeta(shape=None, dtype=None)
# General NumPy array type

FloatArray = ArrayAliasMeta(shape=None, dtype=AnyFloat())
# Arrays containing floating-point numbers

IntArray = ArrayAliasMeta(shape=None, dtype=AnyInt())
# Arrays containing integer numbers

BoolArray = ArrayAliasMeta(shape=None, dtype=np.bool_)
# Arrays containing boolean values

StrArray = ArrayAliasMeta(shape=None, dtype=np.str_)
# Arrays containing string values

Precision-Specific Integer Types

Specific integer precision types for precise memory and performance control.

# Unsigned integer types
ui8 = ArrayAliasMeta(shape=None, dtype=np.uint8)   # 8-bit unsigned integer arrays (0 to 255)
ui16 = ArrayAliasMeta(shape=None, dtype=np.uint16)  # 16-bit unsigned integer arrays (0 to 65,535)
ui32 = ArrayAliasMeta(shape=None, dtype=np.uint32)  # 32-bit unsigned integer arrays (0 to 4.3B)
ui64 = ArrayAliasMeta(shape=None, dtype=np.uint64)  # 64-bit unsigned integer arrays (0 to 18.4E)

# Signed integer types  
i8 = ArrayAliasMeta(shape=None, dtype=np.int8)    # 8-bit signed integer arrays (-128 to 127)
i16 = ArrayAliasMeta(shape=None, dtype=np.int16)   # 16-bit signed integer arrays (-32K to 32K)
i32 = ArrayAliasMeta(shape=None, dtype=np.int32)   # 32-bit signed integer arrays (-2.1B to 2.1B)
i64 = ArrayAliasMeta(shape=None, dtype=np.int64)   # 64-bit signed integer arrays (-9.2E to 9.2E)

Precision-Specific Float Types

Specific floating-point precision types.

f16 = ArrayAliasMeta(shape=None, dtype=np.float16)  # 16-bit half-precision float arrays
f32 = ArrayAliasMeta(shape=None, dtype=np.float32)  # 32-bit single-precision float arrays
f64 = ArrayAliasMeta(shape=None, dtype=np.float64)  # 64-bit double-precision float arrays

Complex Number Types

Complex number array types.

complex64 = ArrayAliasMeta(shape=None, dtype=np.complex64)   # 64-bit complex arrays (32-bit real + 32-bit imag)
complex128 = ArrayAliasMeta(shape=None, dtype=np.complex128)  # 128-bit complex arrays (64-bit real + 64-bit imag)

Special Types

Additional specialized array types.

bool_ = ArrayAliasMeta(shape=None, dtype=np.bool_)  # Boolean array type alias
PRNGKey = ui32[2]  # Random number generator key arrays (2-element uint32 arrays for JAX)

Meta Classes

Metaclass for creating array type aliases.

class ArrayAliasMeta(type):
    """
    Metaclass for creating array type aliases with shape and dtype constraints.
    
    Used internally to create type aliases like f32, i64, etc. that represent
    arrays with specific data types and optional shape constraints.
    """
    def __new__(
        cls, 
        shape: tuple[int, ...] | str | None = None,
        dtype: np.dtype | DType | None = None
    ) -> 'ArrayAliasMeta': ...
    
    def __getitem__(self, shape_spec: str | tuple | int) -> 'ArrayAliasMeta': ...

# DType system classes
class AnyFloat:
    """Matches any floating-point dtype."""
    
class AnyInt:
    """Matches any integer dtype."""
    
class AnyDType:
    """Matches any dtype."""

Usage Examples

Basic Type Annotations

from etils import array_types
import numpy as np

def process_images(
    images: array_types.FloatArray,
    labels: array_types.IntArray
) -> array_types.BoolArray:
    """
    Process images and return prediction correctness.
    
    Args:
        images: Float array of image data
        labels: Integer array of true labels
        
    Returns:
        Boolean array indicating correct predictions
    """
    # Process images...
    predictions = model.predict(images)
    return predictions == labels

def compute_statistics(
    data: array_types.ArrayLike
) -> tuple[array_types.FloatArray, array_types.FloatArray]:
    """
    Compute mean and standard deviation.
    
    Args:
        data: Array-like input data
        
    Returns:
        Tuple of (mean, std) arrays
    """
    arr = np.asarray(data)
    return np.mean(arr), np.std(arr)

Precision-Specific Types

from etils import array_types
import numpy as np

def process_uint8_image(image: array_types.ui8) -> array_types.f32:
    """
    Process 8-bit image and return 32-bit float result.
    
    Args:
        image: 8-bit unsigned integer image (0-255)
        
    Returns:
        32-bit float processed image (0.0-1.0)
    """
    return image.astype(np.float32) / 255.0

def high_precision_computation(
    data: array_types.f64
) -> array_types.f64:
    """
    High-precision numerical computation.
    
    Args:
        data: 64-bit double precision input
        
    Returns:
        64-bit double precision result
    """
    return np.sqrt(data**2 + 1e-15)

def memory_efficient_indices(
    large_array: array_types.Array
) -> array_types.ui32:
    """
    Generate memory-efficient indices.
    
    Args:
        large_array: Input array of any type
        
    Returns:
        32-bit unsigned integer indices
    """
    return np.arange(len(large_array), dtype=np.uint32)

Machine Learning Applications

from etils import array_types
import numpy as np

class NeuralNetwork:
    """Example neural network with typed parameters."""
    
    def __init__(self):
        self.weights: array_types.f32 = np.random.randn(784, 128).astype(np.float32)
        self.biases: array_types.f32 = np.zeros(128, dtype=np.float32)
    
    def forward(
        self,
        inputs: array_types.f32
    ) -> array_types.f32:
        """
        Forward pass through network.
        
        Args:
            inputs: 32-bit float input features
            
        Returns:
            32-bit float network outputs
        """
        return np.dot(inputs, self.weights) + self.biases
    
    def predict(
        self,
        inputs: array_types.f32
    ) -> array_types.IntArray:
        """
        Make predictions with the network.
        
        Args:
            inputs: Input features
            
        Returns:
            Integer class predictions
        """
        logits = self.forward(inputs)
        return np.argmax(logits, axis=-1)

def batch_process(
    batch: array_types.f32,
    targets: array_types.IntArray
) -> dict[str, array_types.FloatArray]:
    """
    Process a batch of training data.
    
    Args:
        batch: Batch of input samples
        targets: Target labels
        
    Returns:
        Dictionary containing loss and accuracy arrays
    """
    model = NeuralNetwork()
    predictions = model.predict(batch)
    
    accuracy = (predictions == targets).mean()
    loss = np.mean((predictions - targets) ** 2)
    
    return {
        'accuracy': np.array([accuracy]),
        'loss': np.array([loss])
    }

Complex Number Operations

from etils import array_types
import numpy as np

def fft_analysis(
    signal: array_types.FloatArray
) -> array_types.complex128:
    """
    Perform FFT analysis on real signal.
    
    Args:
        signal: Real-valued input signal
        
    Returns:
        Complex-valued frequency domain representation
    """
    return np.fft.fft(signal).astype(np.complex128)

def complex_multiplication(
    a: array_types.complex64,
    b: array_types.complex64
) -> array_types.complex64:
    """
    Multiply complex arrays with 64-bit precision.
    
    Args:
        a: First complex array
        b: Second complex array
        
    Returns:
        Product of complex arrays
    """
    return a * b

JAX Integration

from etils import array_types
import numpy as np

def generate_random_key(seed: int = 42) -> array_types.PRNGKey:
    """
    Generate random key for JAX operations.
    
    Args:
        seed: Random seed
        
    Returns:
        PRNG key array for JAX
    """
    # This would typically use jax.random.PRNGKey
    return np.array([seed, 0], dtype=np.uint32)

def random_sampling(
    key: array_types.PRNGKey,
    shape: tuple[int, ...]
) -> array_types.f32:
    """
    Sample random numbers using PRNG key.
    
    Args:
        key: PRNG key for random generation
        shape: Shape of output array
        
    Returns:
        Random float32 array
    """
    # Simplified random generation (would use JAX in practice)
    np.random.seed(int(key[0]))
    return np.random.randn(*shape).astype(np.float32)

String and Boolean Operations

from etils import array_types
import numpy as np

def process_text_data(
    texts: array_types.StrArray
) -> array_types.IntArray:
    """
    Process text data and return length statistics.
    
    Args:
        texts: Array of text strings
        
    Returns:
        Array of text lengths
    """
    return np.array([len(text) for text in texts])

def create_mask(
    data: array_types.FloatArray,
    threshold: float = 0.5
) -> array_types.BoolArray:
    """
    Create boolean mask from float data.
    
    Args:
        data: Input float array
        threshold: Threshold for mask creation
        
    Returns:
        Boolean mask array
    """
    return data > threshold

def filter_by_mask(
    data: array_types.Array,
    mask: array_types.BoolArray
) -> array_types.Array:
    """
    Filter data using boolean mask.
    
    Args:
        data: Input data array
        mask: Boolean mask for filtering
        
    Returns:
        Filtered data array
    """
    return data[mask]

Type Checking and Validation

from etils import array_types
import numpy as np
from typing import TypeGuard

def is_float_array(arr: array_types.Array) -> TypeGuard[array_types.FloatArray]:
    """
    Type guard for float arrays.
    
    Args:
        arr: Array to check
        
    Returns:
        True if array contains float data
    """
    return np.issubdtype(arr.dtype, np.floating)

def validate_array_type(
    arr: array_types.ArrayLike,
    expected_dtype: np.dtype
) -> array_types.Array:
    """
    Validate and convert array to expected type.
    
    Args:
        arr: Input array-like data
        expected_dtype: Expected data type
        
    Returns:
        Validated array with correct dtype
        
    Raises:
        TypeError: If conversion is not possible
    """
    result = np.asarray(arr)
    if result.dtype != expected_dtype:
        try:
            result = result.astype(expected_dtype)
        except (ValueError, TypeError) as e:
            raise TypeError(f"Cannot convert to {expected_dtype}: {e}")
    return result

Install with Tessl CLI