Introducing Action Space

Maintainer's Comment | 17.03.2025

Since I had a major illness, I was unable to post this blog as planned. Also, I wasn't able to code at the pace and with the quality I intended, which led to incomplete code. The reason why I still decided to commit the code was that you guys can see what I've been working on and how I've been progressing. I finally got around to writing about Action Space, a procedural macro designed to compose Tensor operation sequences in Rust. This is still not working properly, but I hope to fix it soon.

Introduction

When working with AI models, especially in reinforcement learning and ensembled learning, you often end up chaining multiple operations on Tensor. Action Space is a Rust procedural macro that makes composing these tensor operations straightforward. With Action Space, the goal was to abstract Tensor operation sequences so that they can be easily composed and reused. This is intended to make it easier to manage and maintain complex Tensor operations.

Example: Composing Tensor Operations

Here's a basic example demonstrating how Action Space helps structure Tensor computations in Rust:

#![allow(unused)]
fn main() {
use cubecl::wgpu::WgpuRuntime;
use lib_proc_macros::action_space;

fn some() {
    let input: TensorHandleRef<WgpuRuntime> = build_tensor::<WgpuRuntime>();

    let result: TensorHandleRef<WgpuRuntime> = action_space!(
        WgpuRuntime,
        (input, ExecMean, output),
        ((input, output), PrepResiduals),
        (ExecSum),
        (PrepSquare),
        (ExecProd),
    );
}
}

Breaking It Down

• The macro takes a Runtime to define the environment.
• The first tuple takes the input before the operations are applied.
• Operations from the PipelineExec and PipelinePush trait are applied sequentially.
• If we want to handle the output of a specific operation differently, we can add a variable name in the tuple behind the operation.
• When the operation needs a non-default input, we can specify it with a tuple before the operation.
• The default case takes the output of the previous operation and passes it to the next operation.

For now, the goal here is that we can input either a 1D, 2D, or 3D Tensor.

• A 1D Tensor results in a Scalar.
• A 2D Tensor results in a Vector.
• A 3D Tensor results in a Matrix.

This isn't working properly yet, and we also do not handle Runtime and Client as intended. This results in the plan to fix and improve this feature next.