mirror of
https://github.com/huggingface/candle.git
synced 2025-06-15 18:28:24 +00:00
e5eb9602d0
* Add support for loading Fortran contiguous tensors This commit introduces the ability to handle Fortran contiguous tensors in the tensor loading process. Previously, the code only supported loading tensors that were contiguous in memory, failing with an error for non-contiguous tensors. With this update, tensors identified as Fortran contiguous (column-major order) are now correctly handled by reversing their dimensions after loading. This enhancement ensures broader compatibility with different tensor layouts, improving the robustness of tensor loading operations. - Check if a tensor is Fortran contiguous using the `is_fortran_contiguous` flag. - For Fortran contiguous tensors, reverse the dimensions after loading to correctly represent their layout in memory. - Continue to bail out with an error for tensors that are neither C contiguous nor Fortran contiguous, maintaining the previous behavior for non-contiguous tensors without explicit support. This change addresses the issue of loading Fortran contiguous tensors, which was previously unsupported, thereby extending the functionality of the tensor loading mechanism to accommodate a wider variety of tensor layouts. * Add reshape step to handle fortran contiguous case * Skip fortran contiguous fix if rank is < 2 * Fail on rank 0, 1 if contiguous
32 lines
919 B
Rust
32 lines
919 B
Rust
/// Regression test for pth files not loading on Windows.
|
|
#[test]
|
|
fn test_pth() {
|
|
let tensors = candle_core::pickle::PthTensors::new("tests/test.pt", None).unwrap();
|
|
tensors.get("test").unwrap().unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn test_pth_with_key() {
|
|
let tensors =
|
|
candle_core::pickle::PthTensors::new("tests/test_with_key.pt", Some("model_state_dict"))
|
|
.unwrap();
|
|
tensors.get("test").unwrap().unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn test_pth_fortran_congiguous() {
|
|
let tensors =
|
|
candle_core::pickle::PthTensors::new("tests/fortran_tensor_3d.pth", None).unwrap();
|
|
let tensor = tensors.get("tensor_fortran").unwrap().unwrap();
|
|
|
|
assert_eq!(tensor.dims3().unwrap(), (2, 3, 4));
|
|
|
|
assert_eq!(
|
|
tensor.to_vec3::<i64>().unwrap(),
|
|
[
|
|
[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
|
|
[[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]
|
|
]
|
|
);
|
|
}
|