I want to write a small neural network in rust from scratch. For best performance, I want to give the rust compiler full knowledge of all vector and matrix sizes and functions used, so it can do as many optimizations as possible. I'm trying to add an activation function as a generic argument to the neural network layer, but it generates a compilation error. Here is what I tried:

我想从头开始写一个小的神经网络。为了获得最佳性能，我希望让RUST编译器充分了解所使用的所有向量和矩阵大小以及函数，这样它就可以进行尽可能多的优化。我试图将激活函数作为泛型参数添加到神经网络层，但它会生成编译错误。以下是我尝试过的：

type VectorFn<const vector_size: usize> = fn(&[f32; vector_size], &mut [f32; vector_size]);

struct Layer<
const input_size: usize,
const output_size: usize,
const activation_fn: VectorFn<output_size>  // error about output_size: the type of const parameters must not depend on other generic parameters
> {
    /// weights
    w: [[f32; input_size]; output_size],

    /// biases
    b: [f32; output_size],
}

// ...

I later want to use the Layer like so:

稍后，我想这样使用该层：

fn VectorizedFn<ScalarFunction>: VectorFn {  # this is pseudo code
    // ...
}
type VectorizedRelu = VectorizedFn<|x| max(x, 0)>  # this is pseudo code

let layer = Layer<32, 64, VectorizedRelu>::new();

I think I have two problems:

1.) There are two cases where I want to hand over a function as a generics argument. The first one is the line where the compiler complains about output_size. The second one is the VectorizedFn<ScalarFunction>. I don't know how to write them.

2.) How to write the functions and "subtypes" of those functions. Do they all have the same type, e.g. something like fn(&[f32; vector_size], &mut [f32; vector_size]) ?

我想我有两个问题：1)在两种情况下，我希望将函数作为泛型参数进行传递。第一行是编译器抱怨输出大小的那一行。第二个是VectorizedFn 。我不知道该怎么写。2.)如何编写函数及其“子类型”。它们是否都具有相同的类型，例如类似FN(&[F32；VECTOR_SIZE]，&MUT[F32；VECTOR_SIZE])的类型？

I'll split this into two parts: taking a function as a generic, and making an alias for a function trait.

我将把它分成两部分：将函数作为泛型，并为函数trait创建别名。

Taking a function as a generic

First, you want a regular generic, not a const generic. I've made your struct simpler so there's less extra info. There are three function traits in Rust: FnOnce, FnMut, and Fn. This uses FnMut.

首先，您希望使用常规泛型，而不是常量泛型。我已经使你的结构更简单了，所以有更少的额外信息。锈菌有三个功能特性：FnOnce、FnMut和Fn。这使用FnMut。

pub struct Layer<const INPUT_SIZE: usize, F> {
    w: [f32; INPUT_SIZE],
    b: [f32; INPUT_SIZE],
    activation: F,
}

impl<const INPUT_SIZE: usize, F> Layer<INPUT_SIZE, F>
where
    F: FnMut(&[f32; INPUT_SIZE], &mut [f32; INPUT_SIZE]),
{
    pub fn activate(&mut self) {
        (self.activation)(&self.b, &mut self.w);
    }
}

Functions (not function pointers) and closures with no captured state have a zero-sized type, so you don't need to worry about activation taking up extra space in your struct unless it needs to, and when you call this, the function will be able to be optimized just like a regular function call.

没有捕获状态的函数(不是函数指针)和闭包具有零大小的类型，因此您不需要担心激活会占用结构中的额外空间，除非它需要，而且当您调用此函数时，函数将能够像常规函数调用一样进行优化。

Trait aliases

In order to avoid writing the function trait signature every time, you can create a "trait alias". This is an informal name for a trait that has no function besides hiding another trait.

为了避免每次都写函数特征签名，可以创建一个“特征别名”。这是对除了隐藏另一个特征之外没有任何作用的特征的非正式名称。

pub trait VectorFn<const VECTOR_SIZE: usize>:
    FnMut(&[f32; VECTOR_SIZE], &mut [f32; VECTOR_SIZE])
{
}

impl<const VECTOR_SIZE: usize, F> VectorFn<VECTOR_SIZE> for F where
    F: FnMut(&[f32; VECTOR_SIZE], &mut [f32; VECTOR_SIZE])
{
}

This consists of two parts. The trait declaration of VectorFn has a const generic, which is then used in its supertrait FnMut. This ensures you can use any VectorFn just like a FnMut. Then, there is a blanket implementation for every F that implements the supertrait. This ensures you can use any FnMut whenever a VectorFn is required.

这包括两个部分。VectorFn的特征声明有一个常量泛型，该泛型随后在其上层特征FnMut中使用。这确保您可以像使用FnMut一样使用任何VectorFn。然后，每个实现上属性的F都有一个完整的实现。这确保无论何时需要VectorFn，您都可以使用任何FnMut。

Note that when making aliases of the function traits, the function arguments must be generics, while the return type should be one or more associated types, in order to match the FnOnce declaration.

请注意，在为函数特征创建别名时，函数参数必须是泛型，而返回类型应该是一个或多个关联类型，以便与FnOnce声明匹配。

Now you can use VectorFn in your type's implementation.

现在，您可以在类型的实现中使用VectorFn。

impl<const INPUT_SIZE: usize, F> Layer<INPUT_SIZE, F>
where
    // F: FnMut(&[f32; INPUT_SIZE], &mut [f32; INPUT_SIZE]),
    F: VectorFn<INPUT_SIZE>,
{
    pub fn activate(&mut self) {
        (self.activation)(&self.b, &mut self.w);
    }
}

Related:

相关：

• How do you pass a Rust function as a parameter?


• How to alias an impl trait?

Here is a complete runnable example for drewtato's answer. This is actually a comment to his answer, but the commenting feature can't take such a long post.

这里有一个完整的可运行的例子来解释德雷瓦茨的答案。这实际上是对他的回答的评论，但评论功能不能接受这么长的帖子。

trait ScalarFn: Fn(f32) -> f32 {}
impl<F> ScalarFn for F where F: Fn(f32) -> f32 {}

pub trait VectorFn<const SIZE: usize>:
    FnMut(&[f32; SIZE], &mut [f32; SIZE]) {}

impl<const SIZE: usize, F> VectorFn<SIZE> for F
where
    F: FnMut(&[f32; SIZE], &mut [f32; SIZE]) {}

pub struct Layer<const SIZE: usize, F> {
    x: [f32; SIZE],
    y: [f32; SIZE],
    activation_fn: F,
}

impl<const SIZE: usize, F>
Layer<SIZE, F>
where
    F: VectorFn<SIZE>,
{
    pub fn new(activation_fn: F) -> Layer<SIZE, F> {
        let x = [0.0_f32; SIZE];
        let y = [0.0_f32; SIZE];

        Layer::<SIZE, F> { x: x, y: y, activation_fn: activation_fn }
    }

    pub fn activate(&mut self) {
        (self.activation_fn)(&self.y, &mut self.x);
    }
}

fn relu_scalar(x: f32) -> f32 {
    0.0_f32.max(x)
}

fn vectorized<const SIZE: usize, F>(inp: &[f32; SIZE], out: &mut [f32; SIZE], scalar_fn: F)
where
    F: ScalarFn,
{
    for i in 0..SIZE {
        out[i] = scalar_fn(inp[i]);
    }
}

fn main() {
    let mut layer = Layer::<10, _>::new(|inp, out| {
        vectorized(inp, out, relu_scalar)
    });

    layer.activate();
}