Struct candle_nn::var_builder::ShardedSafeTensors

pub struct ShardedSafeTensors(/* private fields */);

Implementations

impl ShardedSafeTensors

pub unsafe fn var_builder<P: AsRef<Path>>( paths: &[P], dtype: DType, dev: &Device ) -> Result<ShardedVarBuilder<'static>>

Initializes a VarBuilder that retrieves tensors stored in a collection of safetensors files and make them usable in a sharded way.

Safety

The unsafe is inherited from [memmap2::MmapOptions].

Trait Implementations

impl Backend for ShardedSafeTensors

Get part of a tensor, typically used to do Tensor Parallelism sharding.

If the tensor is of size (1024, 1024).

dim corresponds to the dimension to slice into rank is the rank of the current process world_size is the total number of ranks in the process group

get_sharded("tensor", 0, 0, 2) means tensor.i((..512)) get_sharded("tensor", 0, 1, 2) means tensor.i((512..)) get_sharded("tensor", 1, 0, 2) means tensor.i((.., ..512))

type Hints = Shard

fn get( &self, target_shape: Shape, path: &str, h: Self::Hints, dtype: DType, dev: &Device ) -> Result<Tensor>

Retrieve a tensor with some target shape.

fn contains_tensor(&self, name: &str) -> bool