Skip to main content

Module cache_key

Module cache_key 

Source
Expand description

Cache key trait for flat-index computation.

Built-in implementations: u64, u128, U256, U512, U1024.

§Custom key types

To support index spaces larger than 1024 bits, implement CacheKey for a wider integer type and pass it via CachedFunction::with_key_type:

use bnum::types::U2048;
use tensor4all_tcicore::{CacheKey, CachedFunction};

#[derive(Clone, Hash, PartialEq, Eq)]
struct U2048Key(U2048);

impl CacheKey for U2048Key {
    const BITS_COUNT: u32 = 2048;
    const ZERO: Self = Self(U2048::ZERO);
    const ONE: Self = Self(U2048::ONE);
    fn from_usize(v: usize) -> Self { Self(U2048::from(v as u64)) }
    fn checked_mul(self, rhs: Self) -> Option<Self> {
        self.0.checked_mul(rhs.0).map(Self)
    }
    fn wrapping_add(self, rhs: Self) -> Self {
        Self(self.0.wrapping_add(rhs.0))
    }
}

let local_dims = vec![2usize; 1025];
let cf = CachedFunction::with_key_type::<U2048Key>(
    |idx: &[usize]| idx.iter().sum::<usize>(),
    &local_dims,
).unwrap();
let zeros = vec![0usize; 1025];

assert_eq!(cf.eval(&zeros), 0);
assert_eq!(cf.key_type(), "custom");

Traits§

CacheKey
Trait for cache key types used in flat-index computation.