Merge remote-tracking branch 'origin/main' into cuda-conv-tr1d

candle-core/src/cuda_backend.rs

@@ -608,6 +608,34 @@ impl Map1 for Elu {
     }
 }
 
+struct Col2Im1D {
+    stride: usize,
+}
+
+impl Map1 for Col2Im1D {
+    fn f<T: DeviceRepr + WithDType>(
+        &self,
+        src: &CudaSlice<T>,
+        dev: &CudaDevice,
+        layout: &Layout,
+    ) -> Result<CudaSlice<T>> {
+        let (b_size, l_in, c_out, k_size) = layout.shape().dims4()?;
+        let stride = self.stride;
+        let l_out = (l_in - 1) * stride + k_size;
+
+        let dst_el = b_size * c_out * l_out;
+        let cfg = LaunchConfig::for_num_elems(dst_el as u32);
+        let src = &src.slice(layout.start_offset()..);
+        let func = dev.get_or_load_func(&kernel_name::<T>("col2im1d"), kernels::CONV)?;
+        // SAFETY: Set later by running the kernel.
+        let dst = unsafe { dev.alloc::<T>(dst_el) }.w()?;
+        let params = (l_in, l_out, c_out, k_size, b_size, stride, src, &dst);
+        // SAFETY: ffi.
+        unsafe { func.launch(cfg, params) }.w()?;
+        Ok(dst)
+    }
+}
+
 struct Im2Col1D {
     l_k: usize,
     stride: usize,
@@ -1865,8 +1893,54 @@ impl BackendStorage for CudaStorage {
         params: &crate::conv::ParamsConvTranspose1D,
     ) -> Result<Self> {
         let device = self.device().clone();
-        let slice =
+        const USE_COL2IM_CONV1D_TR: bool = true;
-            ConvTranspose1D(params).map(&self.slice, l, &kernel.slice, kernel_l, &device)?;
+
+        let can_use_col2im = kernel_l.is_contiguous()
+            && params.dilation == 1
+            && params.padding == 0
+            && params.output_padding == 0;
+        if !can_use_col2im || !USE_COL2IM_CONV1D_TR {
+            let slice =
+                ConvTranspose1D(params).map(&self.slice, l, &kernel.slice, kernel_l, &device)?;
+            return Ok(Self { slice, device });
+        }
+
+        let (b_size, c_in, l_in) = l.shape().dims3()?;
+        let (c_in2, c_out, k_size) = kernel_l.shape().dims3()?;
+        if !kernel_l.is_contiguous() {
+            crate::bail!("convtr1d: the second argument (kernel) has to be contiguous {kernel_l:?}")
+        }
+        if c_in != c_in2 {
+            crate::bail!(
+                "convtr1d: shape mismatch on c_in {:?} {:?}",
+                l.shape(),
+                kernel_l.shape()
+            )
+        }
+        let col = {
+            // This merges the last two dimensions of the kernel together.
+            let kernel_l_mm = Layout::new(
+                (b_size, c_in, k_size * c_out).into(),
+                vec![0, k_size * c_out, 1],
+                kernel_l.start_offset(),
+            );
+            self.matmul(
+                kernel,
+                (
+                    b_size,
+                    /* m */ l_in,
+                    /* n */ c_out * k_size,
+                    /* k */ c_in,
+                ),
+                &l.transpose(1, 2)?,
+                &kernel_l_mm,
+            )?
+        };
+        let col_l = Layout::contiguous((b_size, l_in, c_out, k_size));
+        let slice = Col2Im1D {
+            stride: params.stride,
+        }
+        .map(&col.slice, &device, &col_l)?;
         Ok(Self { slice, device })
     }
 

candle-core/src/metal_backend.rs

@@ -609,28 +609,41 @@ impl BackendStorage for MetalStorage {
         let command_buffer = device.command_buffer()?;
         if layout.is_contiguous() && layout.start_offset() == 0 {
             let kernel_name = match (self.dtype, dtype) {
-                (DType::U32, DType::F32) => "cast_u32_f32",
-                (DType::U32, DType::U8) => "cast_u32_u8",
-                (DType::U32, DType::I64) => "cast_u32_i64",
                 (DType::U32, DType::BF16) => "cast_u32_bf16",
+                (DType::U32, DType::F16) => "cast_u32_f16",
+                (DType::U32, DType::F32) => "cast_u32_f32",
+                (DType::U32, DType::I64) => "cast_u32_i64",
+                (DType::U32, DType::U8) => "cast_u32_u8",
 
-                (DType::U8, DType::U32) => "cast_u8_u32",
+                (DType::U8, DType::BF16) => "cast_u8_bf16",
+                (DType::U8, DType::F16) => "cast_u8_f16",
                 (DType::U8, DType::F32) => "cast_u8_f32",
                 (DType::U8, DType::I64) => "cast_u8_i64",
-                (DType::U8, DType::BF16) => "cast_u8_bf16",
+                (DType::U8, DType::U32) => "cast_u8_u32",
 
-                (DType::F32, DType::F16) => "cast_f32_f16",
                 (DType::F32, DType::BF16) => "cast_f32_bf16",
+                (DType::F32, DType::F16) => "cast_f32_f16",
+                (DType::F32, DType::I64) => "cast_f32_i64",
+                (DType::F32, DType::U32) => "cast_f32_u32",
+                (DType::F32, DType::U8) => "cast_f32_u8",
 
+                (DType::I64, DType::BF16) => "cast_i64_bf16",
+                (DType::I64, DType::F16) => "cast_i64_f16",
                 (DType::I64, DType::F32) => "cast_i64_f32",
+                (DType::I64, DType::U32) => "cast_i64_u32",
+                (DType::I64, DType::U8) => "cast_i64_u8",
 
                 (DType::F16, DType::BF16) => "cast_f16_bf16",
                 (DType::F16, DType::F32) => "cast_f16_f32",
+                (DType::F16, DType::I64) => "cast_f16_i64",
+                (DType::F16, DType::U32) => "cast_f16_u32",
+                (DType::F16, DType::U8) => "cast_f16_u8",
 
-                (DType::BF16, DType::U8) => "cast_bf16_u8",
-                (DType::BF16, DType::U32) => "cast_bf16_u32",
                 (DType::BF16, DType::F16) => "cast_bf16_f16",
                 (DType::BF16, DType::F32) => "cast_bf16_f32",
+                (DType::BF16, DType::I64) => "cast_bf16_i64",
+                (DType::BF16, DType::U32) => "cast_bf16_u32",
+                (DType::BF16, DType::U8) => "cast_bf16_u8",
 
                 (left, right) => {
                     crate::bail!("Metal contiguous to_dtype {left:?} {right:?} not implemented")

candle-examples/examples/encodec/main.rs

@@ -109,8 +109,7 @@ fn main() -> Result<()> {
     let codes = match args.action {
         Action::CodeToAudio => {
             let codes = candle::safetensors::load(args.in_file, &device)?;
-            let codes = codes.get("codes").expect("no codes in input file").i(0)?;
+            codes.get("codes").expect("no codes in input file").clone()
-            codes
         }
         Action::AudioToCode | Action::AudioToAudio => {
             let (pcm, sample_rate) = pcm_decode(args.in_file)?;

candle-kernels/src/conv.cu

@@ -51,6 +51,48 @@ __device__ void conv1d(
   dst[dst_i] = static_cast<T>(d);
 }
 
+template <typename T>
+__device__ void col2im1d(
+    const size_t l_in,
+    const size_t l_out,
+    const size_t c_out,
+    const size_t k_size,
+    const size_t b_size,
+    const size_t stride,
+    const T *src,
+    T *dst
+) {
+  const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x;
+  // src: (b_size, l_in, c_out, k_size)
+  // dst: (b_size, c_out, l_out)
+  if (dst_i >= b_size * c_out * l_out) {
+    return;
+  }
+  const size_t dst_s0 = c_out * l_out;
+  const size_t dst_s1 = l_out;
+
+  // dst_idx = b_i * dst_s0 + c_i * dst_s1 + l_in_i * stride + k_i
+  const size_t b_i = dst_i / dst_s0;
+  const size_t dst_i2 = dst_i - b_i * dst_s0;
+  const size_t c_i = dst_i2 / dst_s1;
+  const size_t dst_i3 = dst_i2 - c_i * dst_s1; // l_in_i * stride + k_i
+
+  const size_t src_s0 = c_out * k_size * l_in;
+  const size_t src_s1 = c_out * k_size;
+  const size_t src_s2 = k_size;
+
+  T d = 0;
+  for (size_t k_i = 0; k_i < min(dst_i3 + 1, k_size); ++k_i) {
+    const size_t l_in_i_times_stride = dst_i3 - k_i;
+    const size_t l_in_i = l_in_i_times_stride / stride;
+    const size_t src_i = b_i * src_s0 + l_in_i * src_s1 + c_i * src_s2 + k_i;
+    if (l_in_i * stride == l_in_i_times_stride && l_in_i < l_in) {
+      d += src[src_i];
+    }
+  }
+  dst[dst_i] = d;
+}
+
 template <typename T>
 __device__ void im2col1d(
     const size_t dst_numel,
@@ -527,7 +569,7 @@ extern "C" __global__ void FN_NAME(  \
   conv2d<TYPENAME, TYPEACC>(src_numel, w_out, h_out, stride, padding, dilation, info, src, kernel, dst); \
 } \
 
-#define IM2COL1D_OP(TYPENAME, FN_NAME) \
+#define IM2COL1D_OP(TYPENAME, FN_NAME, FN_NAME2) \
 extern "C" __global__ void FN_NAME(  \
     const size_t dst_numel, \
     const size_t l_out, \
@@ -541,6 +583,18 @@ extern "C" __global__ void FN_NAME(  \
 ) {  \
   im2col1d<TYPENAME>(dst_numel, l_out, l_k, stride, padding, dilation, info, src, dst); \
 } \
+extern "C" __global__ void FN_NAME2(  \
+    const size_t l_in, \
+    const size_t l_out, \
+    const size_t c_out, \
+    const size_t k_size, \
+    const size_t b_size, \
+    const size_t stride, \
+    const TYPENAME *src, \
+    TYPENAME *dst \
+) {  \
+  col2im1d<TYPENAME>(l_in, l_out, c_out, k_size, b_size, stride, src, dst); \
+} \
 
 #define IM2COL_OP(TYPENAME, FN_NAME) \
 extern "C" __global__ void FN_NAME(  \
@@ -642,7 +696,7 @@ AVG_POOL2D_OP(__nv_bfloat16, float, avg_pool2d_bf16)
 MAX_POOL2D_OP(__nv_bfloat16, max_pool2d_bf16)
 UPSAMPLE_NEAREST2D_OP(__nv_bfloat16, upsample_nearest2d_bf16)
 IM2COL_OP(__nv_bfloat16, im2col_bf16)
-IM2COL1D_OP(__nv_bfloat16, im2col1d_bf16)
+IM2COL1D_OP(__nv_bfloat16, im2col1d_bf16, col2im1d_bf16)
 #endif
 
 #if __CUDA_ARCH__ >= 530
@@ -654,7 +708,7 @@ AVG_POOL2D_OP(__half, float, avg_pool2d_f16)
 MAX_POOL2D_OP(__half, max_pool2d_f16)
 UPSAMPLE_NEAREST2D_OP(__half, upsample_nearest2d_f16)
 IM2COL_OP(__half, im2col_f16)
-IM2COL1D_OP(__half, im2col1d_f16)
+IM2COL1D_OP(__half, im2col1d_f16, col2im1d_f16)
 #endif
 
 CONV1D_OP(float, float, conv1d_f32)
@@ -697,7 +751,7 @@ IM2COL_OP(double, im2col_f64)
 IM2COL_OP(uint8_t, im2col_u8)
 IM2COL_OP(uint32_t, im2col_u32)
 
-IM2COL1D_OP(float, im2col1d_f32)
+IM2COL1D_OP(float, im2col1d_f32, col2im1d_f32)
-IM2COL1D_OP(double, im2col1d_f64)
+IM2COL1D_OP(double, im2col1d_f64, col2im1d_f64)
-IM2COL1D_OP(uint8_t, im2col1d_u8)
+IM2COL1D_OP(uint8_t, im2col1d_u8, col2im1d_u8)
-IM2COL1D_OP(uint32_t, im2col1d_u32)
+IM2COL1D_OP(uint32_t, im2col1d_u32, col2im1d_u32)

candle-metal-kernels/src/cast.metal

@@ -72,27 +72,60 @@ kernel void FN_NAME_STRIDED( \
     output[tid] = static_cast<RIGHT_TYPENAME>(static_cast<IR_TYPENAME>(input[get_strided_index(tid, num_dims, dims, strides)])); \
 } \
 
+// u32
 CAST(cast_u32_f32, cast_u32_f32_strided, uint32_t, float)
 CAST(cast_u32_u8, cast_u32_u8_strided, uint32_t, uint8_t)
-CAST(cast_u8_u32, cast_u8_u32_strided, uint8_t, uint32_t)
+CAST(cast_u32_f16, cast_u32_f16_strided, uint32_t, half)
-CAST(cast_u8_f32, cast_u8_f32_strided, uint8_t, float)
-CAST(cast_f16_f32, cast_f16_f32_strided, half, float)
-CAST(cast_f32_f16, cast_f32_f16_strided, float, half)
-
 #if __METAL_VERSION__ >= 220
-CAST(cast_u8_i64, cast_u8_i64_strided, uint8_t, int64_t)
 CAST(cast_u32_i64, cast_u32_i64_strided, uint32_t, int64_t)
-CAST(cast_i64_f32, cast_i64_f32_strided, int64_t, float)
+#endif
+#if defined(__HAVE_BFLOAT__)
+CAST(cast_u32_bf16, cast_u32_bf16_strided, uint32_t, bfloat)
 #endif
 
+// u8
+CAST(cast_u8_u32, cast_u8_u32_strided, uint8_t, uint32_t)
+CAST(cast_u8_f32, cast_u8_f32_strided, uint8_t, float)
+CAST(cast_u8_f16, cast_u8_f16_strided, uint8_t, half)
+#if __METAL_VERSION__ >= 220
+CAST(cast_u8_i64, cast_u8_i64_strided, uint8_t, int64_t)
+#endif
+#if defined(__HAVE_BFLOAT__)
+CAST(cast_u8_bf16, cast_u8_bf16_strided, uint8_t, bfloat)
+#endif
+
+// f16
+CAST(cast_f16_f32, cast_f16_f32_strided, half, float)
+CAST(cast_f16_u8, cast_f16_u8_strided, half, uint8_t)
+CAST(cast_f16_u32, cast_f16_u32_strided, half, uint32_t)
+CAST(cast_f16_i64, cast_f16_i64_strided, half, int64_t)
+#if defined(__HAVE_BFLOAT__)
+CAST_THROUGH(cast_f16_bf16, cast_f16_bf16_strided, half, bfloat, float)
+#endif
+
+// i64
+CAST(cast_i64_f32, cast_i64_f32_strided, int64_t, float)
+CAST(cast_i64_u8, cast_i64_u8_strided, int64_t, uint8_t)
+CAST(cast_i64_u32, cast_i64_u32_strided, int64_t, uint32_t)
+CAST(cast_i64_f16, cast_i64_f16_strided, int64_t, half)
+#if defined(__HAVE_BFLOAT__)
+CAST_THROUGH(cast_i64_bf16, cast_i64_bf16_strided, int64_t, bfloat, float)
+#endif
+
+// f32
+CAST(cast_f32_f16, cast_f32_f16_strided, float, half)
+CAST(cast_f32_u32, cast_f32_u32_strided, float, uint32_t)
+CAST(cast_f32_u8, cast_f32_u8_strided, float, uint8_t)
+CAST(cast_f32_i64, cast_f32_i64_strided, float, int64_t)
+#if defined(__HAVE_BFLOAT__)
+CAST(cast_f32_bf16, cast_f32_bf16_strided, float, bfloat)
+#endif
+
+// bf16
 #if defined(__HAVE_BFLOAT__)
 CAST(cast_bf16_u32, cast_bf16_u32_strided, bfloat, uint32_t)
+CAST(cast_bf16_i64, cast_bf16_i64_strided, bfloat, int64_t)
 CAST(cast_bf16_f32, cast_bf16_f32_strided, bfloat, float)
-CAST(cast_u8_bf16, cast_u8_bf16_strided, uint8_t, bfloat)
-CAST(cast_u32_bf16, cast_u32_bf16_strided, uint32_t, bfloat)
-CAST(cast_f32_bf16, cast_f32_bf16_strided, float, bfloat)
-
 CAST_THROUGH(cast_bf16_u8, cast_bf16_u8_strided, bfloat, uint8_t, float)
 CAST_THROUGH(cast_bf16_f16, cast_bf16_f16_strided, bfloat, half, float)
-CAST_THROUGH(cast_f16_bf16, cast_f16_bf16_strided, half, bfloat, float)
 #endif

candle-metal-kernels/src/tests.rs

@@ -292,7 +292,7 @@ fn binary_ops_bf16() {
     binary_op!(max, |x: bf16, y| x.max(y));
 }
 
-fn cast<T: Clone, U: Clone>(v: &[T], name: &'static str) -> Vec<U> {
+fn run_cast<T: Clone, U: Clone>(v: &[T], name: &'static str) -> Vec<U> {
     let device = device();
     let kernels = Kernels::new();
     let command_queue = device.new_command_queue();
@@ -319,107 +319,189 @@ fn cast<T: Clone, U: Clone>(v: &[T], name: &'static str) -> Vec<U> {
 }
 
 #[test]
-fn cast_u32_f32() {
+fn cast_f32() {
-    let v = vec![1u32, 2, 3];
+    let v_f64 = vec![1.0f64, 2.0, 3.0];
-    let results = cast(&v, "cast_u32_f32");
+    let v_f32: Vec<f32> = v_f64.iter().map(|&v| v as f32).collect();
-    let expected: Vec<_> = v.iter().map(|&v| v as f32).collect();
+    let v_f16: Vec<f16> = v_f64.iter().map(|&v| f16::from_f32(v as f32)).collect();
-    assert_eq!(approx(results, 4), vec![1.0f32, 2.0, 3.0]);
+    let v_bf16: Vec<bf16> = v_f64.iter().map(|&v| bf16::from_f32(v as f32)).collect();
-    assert_eq!(approx(expected, 4), vec![1.0f32, 2.0, 3.0]);
+    let v_u32: Vec<u32> = v_f64.iter().map(|&v| v as u32).collect();
+    let v_u8: Vec<u8> = v_f64.iter().map(|&v| v as u8).collect();
+    let v_i64: Vec<i64> = v_f64.iter().map(|&v| v as i64).collect();
 
-    let v = vec![1.0f32, 2.0, 3.0];
+    // f32 -> f16
-    let input: Vec<f16> = v.iter().map(|v| f16::from_f32(*v)).collect();
+    let results: Vec<half::f16> = run_cast(&v_f32, "cast_f32_f16");
-    let results: Vec<f32> = cast(&input, "cast_f16_f32");
+    assert_eq!(results, v_f16);
-    assert_eq!(results, vec![1.0f32, 2.0, 3.0]);
 
-    let v = vec![1.0f32; 10_000];
+    // f32 -> bf16
-    let input: Vec<f16> = v.iter().map(|v| f16::from_f32(*v)).collect();
+    let results: Vec<bf16> = run_cast(&v_f32, "cast_f32_bf16");
-    let results: Vec<f32> = cast(&input, "cast_f16_f32");
+    assert_eq!(results, v_bf16);
-    assert_eq!(results.len(), 10_000);
+
-    assert_eq!(&results[..10], vec![1.0f32; 10]);
+    // f32 -> u32
-    assert_eq!(results, vec![1.0f32; 10_000]);
+    let results: Vec<u32> = run_cast(&v_f32, "cast_f32_u32");
+    assert_eq!(results, v_u32);
+
+    // f32 -> u8
+    let results: Vec<u8> = run_cast(&v_f32, "cast_f32_u8");
+    assert_eq!(results, v_u8);
+
+    // f32 -> i64
+    let results: Vec<i64> = run_cast(&v_f32, "cast_f32_i64");
+    assert_eq!(results, v_i64);
 }
 
 #[test]
-fn it_cast_bf16_u32() {
+fn cast_f16() {
-    let input: Vec<bf16> = (1..=3).map(|v| bf16::from_f32(v as f32)).collect();
+    let v_f64 = vec![1.0f64, 2.0, 3.0];
+    let v_f32: Vec<f32> = v_f64.iter().map(|&v| v as f32).collect();
+    let v_f16: Vec<f16> = v_f64.iter().map(|&v| f16::from_f32(v as f32)).collect();
+    let v_bf16: Vec<bf16> = v_f64.iter().map(|&v| bf16::from_f32(v as f32)).collect();
+    let v_u32: Vec<u32> = v_f64.iter().map(|&v| v as u32).collect();
+    let v_u8: Vec<u8> = v_f64.iter().map(|&v| v as u8).collect();
+    let v_i64: Vec<i64> = v_f64.iter().map(|&v| v as i64).collect();
 
-    let output: Vec<u32> = cast(&input, "cast_bf16_u32");
+    // f16 -> f32
-    let expected: Vec<u32> = (1..=3).map(|v| v as u32).collect();
+    let results: Vec<f32> = run_cast(&v_f16, "cast_f16_f32");
+    assert_eq!(results, v_f32);
 
-    assert_eq!(output, expected);
+    // f16 -> bf16
+    let results: Vec<bf16> = run_cast(&v_f16, "cast_f16_bf16");
+    assert_eq!(results, v_bf16);
+
+    // f16 -> u32
+    let results: Vec<u32> = run_cast(&v_f16, "cast_f16_u32");
+    assert_eq!(results, v_u32);
+
+    // f16 -> u8
+    let results: Vec<u8> = run_cast(&v_f16, "cast_f16_u8");
+    assert_eq!(results, v_u8);
+
+    // f16 -> i64
+    let results: Vec<i64> = run_cast(&v_f16, "cast_f16_i64");
+    assert_eq!(results, v_i64);
 }
 
 #[test]
-fn it_cast_bf16_f32() {
+fn cast_bf16() {
-    let input: Vec<bf16> = (1..=3).map(|v| bf16::from_f32(v as f32)).collect();
+    let v_f64 = vec![1.0f64, 2.0, 3.0];
+    let v_f32: Vec<f32> = v_f64.iter().map(|&v| v as f32).collect();
+    let v_f16: Vec<f16> = v_f64.iter().map(|&v| f16::from_f32(v as f32)).collect();
+    let v_bf16: Vec<bf16> = v_f64.iter().map(|&v| bf16::from_f32(v as f32)).collect();
+    let v_u32: Vec<u32> = v_f64.iter().map(|&v| v as u32).collect();
+    let v_u8: Vec<u8> = v_f64.iter().map(|&v| v as u8).collect();
+    let v_i64: Vec<i64> = v_f64.iter().map(|&v| v as i64).collect();
 
-    let output: Vec<f32> = cast(&input, "cast_bf16_f32");
+    // bf16 -> f32
-    let expected: Vec<f32> = (1..=3).map(|v| v as f32).collect();
+    let results: Vec<f32> = run_cast(&v_bf16, "cast_bf16_f32");
+    assert_eq!(results, v_f32);
 
-    assert_eq!(output, expected);
+    // bf16 -> f16
+    let results: Vec<f16> = run_cast(&v_bf16, "cast_bf16_f16");
+    assert_eq!(results, v_f16);
+
+    // bf16 -> u32
+    let results: Vec<u32> = run_cast(&v_bf16, "cast_bf16_u32");
+    assert_eq!(results, v_u32);
+
+    // bf16 -> u8
+    let results: Vec<u8> = run_cast(&v_bf16, "cast_bf16_u8");
+    assert_eq!(results, v_u8);
+
+    // bf16 -> i64
+    let results: Vec<i64> = run_cast(&v_bf16, "cast_bf16_i64");
+    assert_eq!(results, v_i64);
 }
 
 #[test]
-fn it_cast_u8_bf16() {
+fn cast_u32() {
-    let input: Vec<u8> = (1..=3).map(|v| v as u8).collect();
+    let v_f64 = vec![1.0f64, 2.0, 3.0];
+    let v_f32: Vec<f32> = v_f64.iter().map(|&v| v as f32).collect();
+    let v_f16: Vec<f16> = v_f64.iter().map(|&v| f16::from_f32(v as f32)).collect();
+    let v_bf16: Vec<bf16> = v_f64.iter().map(|&v| bf16::from_f32(v as f32)).collect();
+    let v_u32: Vec<u32> = v_f64.iter().map(|&v| v as u32).collect();
+    let v_u8: Vec<u8> = v_f64.iter().map(|&v| v as u8).collect();
+    let v_i64: Vec<i64> = v_f64.iter().map(|&v| v as i64).collect();
 
-    let output: Vec<bf16> = cast(&input, "cast_u8_bf16");
+    // u32 -> f32
-    let expected: Vec<bf16> = input
+    let results: Vec<f32> = run_cast(&v_u32, "cast_u32_f32");
-        .iter()
+    assert_eq!(results, v_f32);
-        .map(|v| bf16::from_f32(*v as f32))
-        .collect::<Vec<_>>();
 
-    assert_eq!(output, expected);
+    // u32 -> f16
+    let results: Vec<f16> = run_cast(&v_u32, "cast_u32_f16");
+    assert_eq!(results, v_f16);
+
+    // u32 -> bf16
+    let results: Vec<bf16> = run_cast(&v_u32, "cast_u32_bf16");
+    assert_eq!(results, v_bf16);
+
+    // u32 -> u8
+    let results: Vec<u8> = run_cast(&v_u32, "cast_u32_u8");
+    assert_eq!(results, v_u8);
+
+    // u32 -> i64
+    let results: Vec<i64> = run_cast(&v_u32, "cast_u32_i64");
+    assert_eq!(results, v_i64);
 }
 
 #[test]
-fn it_cast_u32_bf16() {
+fn cast_u8() {
-    let input: Vec<u32> = (1..=3).map(|v| v as u32).collect();
+    let v_f64 = vec![1.0f64, 2.0, 3.0];
+    let v_f32: Vec<f32> = v_f64.iter().map(|&v| v as f32).collect();
+    let v_f16: Vec<f16> = v_f64.iter().map(|&v| f16::from_f32(v as f32)).collect();
+    let v_bf16: Vec<bf16> = v_f64.iter().map(|&v| bf16::from_f32(v as f32)).collect();
+    let v_u32: Vec<u32> = v_f64.iter().map(|&v| v as u32).collect();
+    let v_u8: Vec<u8> = v_f64.iter().map(|&v| v as u8).collect();
+    let v_i64: Vec<i64> = v_f64.iter().map(|&v| v as i64).collect();
 
-    let output: Vec<bf16> = cast(&input, "cast_u32_bf16");
+    // u8 -> f32
-    let expected: Vec<bf16> = input.iter().map(|v| bf16::from_f32(*v as f32)).collect();
+    let results: Vec<f32> = run_cast(&v_u8, "cast_u8_f32");
+    assert_eq!(results, v_f32);
 
-    assert_eq!(output, expected);
+    // u8 -> f16
+    let results: Vec<f16> = run_cast(&v_u8, "cast_u8_f16");
+    assert_eq!(results, v_f16);
+
+    // u8 -> bf16
+    let results: Vec<bf16> = run_cast(&v_u8, "cast_u8_bf16");
+    assert_eq!(results, v_bf16);
+
+    // u8 -> u32
+    let results: Vec<u32> = run_cast(&v_u8, "cast_u8_u32");
+    assert_eq!(results, v_u32);
+
+    // u8 -> i64
+    let results: Vec<i64> = run_cast(&v_u8, "cast_u8_i64");
+    assert_eq!(results, v_i64);
 }
 
 #[test]
-fn it_cast_f32_bf16() {
+fn cast_i64() {
-    let input: Vec<f32> = (1..=3).map(|v| v as f32).collect();
+    let v_f64 = vec![1.0f64, 2.0, 3.0];
+    let v_f32: Vec<f32> = v_f64.iter().map(|&v| v as f32).collect();
+    let v_f16: Vec<f16> = v_f64.iter().map(|&v| f16::from_f32(v as f32)).collect();
+    let v_bf16: Vec<bf16> = v_f64.iter().map(|&v| bf16::from_f32(v as f32)).collect();
+    let v_u32: Vec<u32> = v_f64.iter().map(|&v| v as u32).collect();
+    let v_u8: Vec<u8> = v_f64.iter().map(|&v| v as u8).collect();
+    let v_i64: Vec<i64> = v_f64.iter().map(|&v| v as i64).collect();
 
-    let output: Vec<bf16> = cast(&input, "cast_f32_bf16");
+    // i64 -> f32
-    let expected: Vec<bf16> = input.iter().map(|v| bf16::from_f32(*v as f32)).collect();
+    let results: Vec<f32> = run_cast(&v_i64, "cast_i64_f32");
+    assert_eq!(results, v_f32);
 
-    assert_eq!(output, expected);
+    // i64 -> f16
-}
+    let results: Vec<f16> = run_cast(&v_i64, "cast_i64_f16");
+    assert_eq!(results, v_f16);
 
-#[test]
+    // i64 -> bf16
-fn it_cast_bf16_u8() {
+    let results: Vec<bf16> = run_cast(&v_i64, "cast_i64_bf16");
-    let input: Vec<bf16> = (1..=3).map(|v| bf16::from_f32(v as f32)).collect();
+    assert_eq!(results, v_bf16);
 
-    let output: Vec<u8> = cast(&input, "cast_bf16_u8");
+    // i64 -> u32
-    let expected: Vec<u8> = input.iter().map(|v| v.to_f32() as u8).collect();
+    let results: Vec<u32> = run_cast(&v_i64, "cast_i64_u32");
+    assert_eq!(results, v_u32);
 
-    assert_eq!(output, expected);
+    // i64 -> u8
-}
+    let results: Vec<u8> = run_cast(&v_i64, "cast_i64_u8");
-
+    assert_eq!(results, v_u8);
-#[test]
-fn it_cast_bf16_f16() {
-    let input: Vec<bf16> = (1..=3).map(|v| bf16::from_f32(v as f32)).collect();
-
-    let output: Vec<f16> = cast(&input, "cast_bf16_f16");
-    let expected: Vec<f16> = input.iter().map(|v| f16::from_f32(v.to_f32())).collect();
-
-    assert_eq!(output, expected);
-}
-
-#[test]
-fn it_cast_f16_bf16() {
-    let input: Vec<f16> = (1..=3).map(|v| f16::from_f32(v as f32)).collect();
-
-    let output: Vec<bf16> = cast(&input, "cast_f16_bf16");
-    let expected: Vec<bf16> = input.iter().map(|v| bf16::from_f32(v.to_f32())).collect();
-
-    assert_eq!(output, expected);
 }
 
 fn run_affine<T: Clone>(v: &[T], mul: f64, add: f64) -> Vec<T> {

candle-transformers/src/models/mod.rs

@@ -23,7 +23,6 @@ pub mod mistral;
 pub mod mixformer;
 pub mod mixtral;
 pub mod mobileone;
-pub mod moondream;
 pub mod mpt;
 pub mod persimmon;
 pub mod phi;

candle-transformers/src/models/moondream.rs

@@ -1,174 +0,0 @@
-#![allow(unused)]
-use crate::models::phi;
-use candle::{Module, Result, Tensor};
-use candle_nn::{linear_b, Linear, VarBuilder};
-
-// https://github.com/vikhyat/moondream/blob/main/moondream/configuration_moondream.py
-#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
-pub struct Config {
-    phi_config: phi::Config,
-    vision_config: VisionConfig,
-}
-
-#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
-pub struct VisionConfig {
-    image_embedding_dim: usize,
-    model_dim: usize,
-    hidden_dim: usize,
-    act: candle_nn::Activation,
-}
-
-impl VisionConfig {
-    pub fn v2() -> Self {
-        Self {
-            image_embedding_dim: 1152,
-            model_dim: 2048,
-            hidden_dim: 2048 * 4,
-            act: candle_nn::Activation::Silu,
-        }
-    }
-}
-
-impl Config {
-    pub fn v2() -> Self {
-        let phi_config = phi::Config {
-            vocab_size: 51200,
-            hidden_size: 2048,
-            intermediate_size: 8192,
-            num_hidden_layers: 24,
-            num_attention_heads: 32,
-            num_key_value_heads: None,
-            hidden_act: candle_nn::Activation::NewGelu,
-            max_position_embeddings: 2048,
-            tie_word_embeddings: false,
-            layer_norm_eps: 1e-5,
-            rope_theta: 10_000.,
-            partial_rotary_factor: 0.5,
-            qk_layernorm: false,
-        };
-        let vision_config = VisionConfig::v2();
-        Self {
-            phi_config,
-            vision_config,
-        }
-    }
-}
-
-#[derive(Debug, Clone)]
-struct LinearPatchEmbedding {
-    linear: Linear,
-}
-
-#[derive(Debug, Clone)]
-struct Encoder {}
-
-impl Encoder {
-    fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result<Self> {
-        todo!()
-    }
-}
-
-impl Module for Encoder {
-    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        todo!()
-    }
-}
-
-#[derive(Debug, Clone)]
-struct Mlp {
-    fc1: Linear,
-    act: candle_nn::Activation,
-    fc2: Linear,
-}
-
-impl Mlp {
-    fn new(
-        in_f: usize,
-        hidden_f: usize,
-        out_f: usize,
-        act: candle_nn::Activation,
-        vb: VarBuilder,
-    ) -> Result<Self> {
-        let fc1 = linear_b(in_f, hidden_f, true, vb.pp("fc1"))?;
-        let fc2 = linear_b(hidden_f, out_f, true, vb.pp("fc2"))?;
-        Ok(Self { fc1, act, fc2 })
-    }
-}
-
-impl Module for Mlp {
-    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        xs.apply(&self.fc1)?.apply(&self.act)?.apply(&self.fc2)
-    }
-}
-
-#[derive(Debug, Clone)]
-struct VisionProjection {
-    mlp: Mlp,
-}
-
-impl VisionProjection {
-    fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result<Self> {
-        let mlp = Mlp::new(
-            cfg.image_embedding_dim,
-            cfg.hidden_dim,
-            cfg.model_dim,
-            cfg.act,
-            vb.pp("mlp"),
-        )?;
-        Ok(Self { mlp })
-    }
-}
-
-impl Module for VisionProjection {
-    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        xs.apply(&self.mlp)
-    }
-}
-
-#[derive(Debug, Clone)]
-struct VisionEncoder {
-    encoder: Encoder,
-    projection: VisionProjection,
-}
-
-impl VisionEncoder {
-    pub fn new(cfg: &VisionConfig, vb: VarBuilder) -> Result<Self> {
-        let encoder = Encoder::new(cfg, vb.pp("vision.trunk"))?;
-        let projection = VisionProjection::new(cfg, vb.pp("projection"))?;
-        Ok(Self {
-            encoder,
-            projection,
-        })
-    }
-}
-
-impl Module for VisionEncoder {
-    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        let (b, c, hp1, wp2) = xs.dims4()?;
-        let (p1, p2) = (14, 14);
-        let h = hp1 / p1;
-        let w = wp2 / p2;
-        let xs = xs
-            .reshape((b, c, h, p1, h, p2))?
-            .permute((0, 2, 4, 1, 3, 5))?
-            .reshape((b, h * w, c * p1 * p2))?;
-        xs.apply(&self.encoder)?.apply(&self.projection)
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct Model {
-    text_model: phi::Model,
-    vision_encoder: VisionEncoder,
-}
-
-impl Model {
-    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
-        let text_model = phi::Model::new(&cfg.phi_config, vb.pp("text_model"))?;
-        let vision_encoder = VisionEncoder::new(&cfg.vision_config, vb.pp("vision_encoder"))?;
-        Ok(Self {
-            text_model,
-            vision_encoder,
-        })
-    }
-}

candle-transformers/src/models/phi.rs

@@ -106,7 +106,7 @@ impl Module for MLP {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 struct Attention {
     q_proj: Linear,
     k_proj: Linear,
@@ -265,7 +265,7 @@ impl Attention {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 struct DecoderLayer {
     self_attn: Attention,
     mlp: MLP,
@@ -304,7 +304,7 @@ impl DecoderLayer {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct Model {
     embed_tokens: Embedding,
     layers: Vec<DecoderLayer>,