burn_jit/kernel/
contiguous.rs

1use crate::{execute_with_dtype, tensor::JitTensor, JitRuntime};
2
3/// Make a jit tensor contiguous.
4pub fn into_contiguous<R: JitRuntime>(tensor: JitTensor<R>) -> JitTensor<R> {
5    if tensor.is_contiguous() {
6        return tensor;
7    }
8
9    execute_with_dtype!(tensor.dtype, E, {
10        let output = cubecl::linalg::tensor::into_contiguous::<R, E>(
11            &tensor.client,
12            &tensor.as_handle_ref(),
13        );
14
15        JitTensor::new(
16            tensor.client,
17            output.handle,
18            output.shape.into(),
19            tensor.device,
20            output.strides,
21            tensor.dtype,
22        )
23    })
24}