torsh-python 0.1.3

//! Python bindings for ToRSh - PyTorch-compatible deep learning in Rust
//!
//! This crate provides Python bindings for the ToRSh deep learning framework,
//! enabling PyTorch-compatible APIs to be used from Python.
//!
//! # Modular Structure
//!
//! The crate is organized into focused modules:
//! - `tensor` - Tensor operations and creation functions
//! - `nn` - Neural network layers and containers
//! - `optim` - Optimization algorithms
//! - `device` - Device management and utilities
//! - `dtype` - Data type definitions and conversions
//! - `error` - Error handling and conversions
//! - `utils` - Common utilities and helpers

use pyo3::prelude::*;

// Core modules - modular structure
pub mod device;
pub mod dtype;
pub mod error;
pub mod nn;
pub mod optim;
pub mod tensor;
pub mod utils;

pub mod autograd;
pub mod data;
pub mod distributed;
// pub mod functional;  // Disabled until tensor ops are fully implemented

// Re-export main types
pub use device::PyDevice;
pub use dtype::PyDType;
pub use error::TorshPyError;
pub use tensor::PyTensor;

/// ToRSh Python module
#[pymodule]
fn rstorch(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Register main classes
    m.add_class::<PyTensor>()?;
    m.add_class::<PyDevice>()?;
    m.add_class::<PyDType>()?;

    // Add submodules with new modular structure
    nn::register_nn_module(m.py(), m)?;
    optim::register_optim_module(m.py(), m)?;

    let data_module = PyModule::new(m.py(), "data")?;
    data::register_data_module(m.py(), &data_module)?;
    m.add_submodule(&data_module)?;

    let autograd_module = PyModule::new(m.py(), "autograd")?;
    autograd::register_autograd_module(m.py(), &autograd_module)?;
    m.add_submodule(&autograd_module)?;

    let distributed_module = PyModule::new(m.py(), "distributed")?;
    distributed::register_distributed_module(m.py(), &distributed_module)?;
    m.add_submodule(&distributed_module)?;

    // let functional_module = PyModule::new(m.py(), "F")?;
    // functional::register_functional_module(m.py(), &functional_module)?;
    // m.add_submodule(&functional_module)?;

    // Add tensor creation functions
    tensor::register_creation_functions(m)?;

    // Add device and dtype constants
    device::register_device_constants(m)?;
    dtype::register_dtype_constants(m)?;

    // Register error types
    error::register_error_types(m)?;

    // Set version
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;

    Ok(())
}

/// A Python module implemented in Rust.
#[pymodule]
fn rstorch_python(m: &Bound<'_, PyModule>) -> PyResult<()> {
    rstorch(m)
}