use tenferro_tensor::{DType, Tensor, TypedTensor};
use crate::error::Result;
pub use crate::checkpoint::RetainedValue;
pub use crate::metadata::tensor_meta_from_tensor;
pub fn next_input_key() -> tenferro_ops::input_key::TensorInputKey {
crate::traced::next_input_key()
}
pub fn ones_tensor(dtype: DType, shape: Vec<usize>) -> Result<Tensor> {
match dtype {
DType::F32 => Ok(Tensor::from_typed::<f32>(TypedTensor::ones(shape)?)),
DType::F64 => Ok(Tensor::from_typed::<f64>(TypedTensor::ones(shape)?)),
DType::I32 => Ok(Tensor::from_typed::<i32>(TypedTensor::ones(shape)?)),
DType::I64 => Ok(Tensor::from_typed::<i64>(TypedTensor::ones(shape)?)),
DType::Bool => {
let len =
tenferro_tensor::validate::checked_shape_product("ones_tensor", "shape", &shape)?;
Ok(Tensor::from_typed::<bool>(TypedTensor::from_vec_col_major(
shape,
vec![true; len],
)?))
}
DType::C32 => Ok(Tensor::from_typed::<tenferro_tensor::Complex32>(
TypedTensor::ones(shape)?,
)),
DType::C64 => Ok(Tensor::from_typed::<tenferro_tensor::Complex64>(
TypedTensor::ones(shape)?,
)),
DType::External(_) => Err(crate::Error::invalid_argument(
"ones_tensor",
crate::ErrorPhase::GraphBuild,
"dtype",
"an externally defined scalar has no core runtime tensor",
)),
}
}
#[cfg(test)]
mod tests;