types

Type definitions for stochpw.

This module provides type aliases for JAX-specific types to improve type safety and reduce reliance on Any types throughout the codebase.