axonml-core

Overview

axonml-core is the foundational layer of the AxonML machine learning framework. It provides core abstractions for device management, memory storage, data types, and pluggable backend implementations that underpin all tensor operations across CPU and GPU devices.

Features

• Device Abstraction - Unified interface for managing compute devices including CPU, CUDA, Vulkan, Metal, and WebGPU backends with seamless tensor transfer between devices.

• Type-Safe Data Types - Comprehensive type system supporting f16, f32, f64, i8, i16, i32, i64, u8, u32, u64, and bool with automatic type promotion rules.

• Efficient Memory Storage - Reference-counted storage with zero-copy slicing, automatic memory cleanup, and device-agnostic operations.

• Pluggable Backend Architecture - Extensible backend system with a common Backend trait enabling device-agnostic tensor operations.

• Memory Allocator - Flexible allocator trait with default CPU implementation and support for custom memory pools.

• Device Capabilities - Query device capabilities including memory, f16/f64 support, and compute capability for optimal resource utilization.

Modules

Usage

Add this to your Cargo.toml:

[dependencies]
axonml-core = "0.1.0"

Basic Example

use axonml_core::{Device, DType, Storage};

// Check device availability
let device = Device::Cpu;
assert!(device.is_available());

// Create storage on CPU
let storage = Storage::<f32>::zeros(1024, device);
assert_eq!(storage.len(), 1024);

// Create storage from data
let data = vec![1.0f32, 2.0, 3.0, 4.0];
let storage = Storage::from_vec(data, Device::Cpu);

// Create a view (zero-copy slice)
let view = storage.slice(1, 2).unwrap();
assert_eq!(view.len(), 2);

Device Capabilities

use axonml_core::Device;

let device = Device::Cpu;
let caps = device.capabilities();

println!("Device: {}", caps.name);
println!("Total Memory: {} bytes", caps.total_memory);
println!("Supports f16: {}", caps.supports_f16);
println!("Supports f64: {}", caps.supports_f64);

Data Types

use axonml_core::{DType, Scalar, Numeric, Float};

// Query dtype properties
assert!(DType::F32.is_float());
assert_eq!(DType::F32.size_of(), 4);

// Use type traits
fn process<T: Float>(data: &[T]) -> T {
    data.iter().fold(T::ZERO, |acc, &x| acc + x)
}

Tests

Run the test suite:

cargo test -p axonml-core

License

Licensed under either of:

• Apache License, Version 2.0 (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
• MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT)

at your option.