candle/candle-examples/examples/custom-ops

candle-custom-ops

This example illustrates how to implement forward and backward passes for custom operations on the CPU and GPU. The custom op in this example implements RMS normalization for the CPU and CUDA.

Running an example

$ cargo run --example custom-ops

> [[ 0.,  1.,  2.,  3.,  4.,  5.,  6.],
>  [ 7.,  8.,  9., 10., 11., 12., 13.]]
> Tensor[[2, 7], f32]
> [[0.0000, 0.2773, 0.5547, 0.8320, 1.1094, 1.3867, 1.6641],
>  [0.6864, 0.7845, 0.8825, 0.9806, 1.0786, 1.1767, 1.2748]]
> Tensor[[2, 7], f32]