Bimm Shape Contracts

• static/stack-evaluated runtime shape contracts for tensors.

Example Usage

/// Window Partition
///
/// ## Parameters
///
/// - `tensor`: Input tensor of shape (B, h_wins * window_size, w_wins * window_size, C).
/// - `window_size`: Window size.
///
/// ## Returns
/// 
/// Output tensor of shape (B * h_windows * w_windows, window_size, window_size, C).
/// 
/// ## Panics
/// 
/// Panics if the input tensor does not have 4 dimensions.
pub fn window_partition<B: Backend, K>(
    tensor: Tensor<B, 4, K>,
    window_size: usize,
) -> Tensor<B, 4, K>
where
    K: BasicOps<B>,
{
    static CONTRACT: ShapeContract = ShapeContract::new(&[
        DimMatcher::Expr(DimExpr::Param("batch")),
        DimMatcher::Expr(DimExpr::Prod(&[
            DimExpr::Param("h_wins"),
            DimExpr::Param("window_size"),
        ])),
        DimMatcher::Expr(DimExpr::Prod(&[
            DimExpr::Param("w_wins"),
            DimExpr::Param("window_size"),
        ])),
        DimMatcher::Expr(DimExpr::Param("channels")),
    ]);
    let [b, h_wins, w_wins, c] = CONTRACT.unpack_shape(
        &tensor.dims(),
        &["batch", "h_wins", "w_wins", "channels"],
        &[("window_size", window_size)],
    );

    tensor
        .reshape([b, h_wins, window_size, w_wins, window_size, c])
        .swap_dims(2, 3)
        .reshape([b * h_wins * w_wins, window_size, window_size, c])
}

Performance

Benchmark: 214.11 ns/iter (+/- 4.71)

#[bench]
fn bench_shape_contract(b: &mut Bencher) {
    static PATTERN: ShapeContract = ShapeContract::new(&[
        DimMatcher::Any,
        DimMatcher::Expr(DimExpr::Param("b")),
        DimMatcher::Ellipsis,
        DimMatcher::Expr(DimExpr::Prod(&[DimExpr::Param("h"), DimExpr::Param("p")])),
        DimMatcher::Expr(DimExpr::Prod(&[DimExpr::Param("w"), DimExpr::Param("p")])),
        DimMatcher::Expr(DimExpr::Pow(&DimExpr::Param("z"), 3)),
        DimMatcher::Expr(DimExpr::Param("c")),
    ]);

    let batch = 2;
    let height = 3;
    let width = 2;
    let padding = 4;
    let channels = 5;
    let z = 4;

    let shape = [12, batch, 1, 2, 3, height * padding, width * padding, z * z * z, channels];
    let env = [("p", padding), ("c", channels)];
    let keys = ["b", "h", "w", "z"];

    b.iter(|| {
        let _ = PATTERN.unpack_shape(&shape, &keys, &env);
    });
}

assert_shape_every_n

There is also the assert_shape_every_n method that that can be used to assert the shape of a tensor at runtime, but only every n calls, starting with the first.

Note that this method does not return keys, but only checks the shape.

Benchmark: 20.66 ns/iter (+/- 0.51)

#[bench]
fn bench_assert_shape_every_n(b: &mut Bencher) {
    static PATTERN: ShapeContract = ShapeContract::new(&[
        DimMatcher::Any,
        DimMatcher::Expr(DimExpr::Param("b")),
        DimMatcher::Ellipsis,
        DimMatcher::Expr(DimExpr::Prod(&[DimExpr::Param("h"), DimExpr::Param("p")])),
        DimMatcher::Expr(DimExpr::Prod(&[DimExpr::Param("w"), DimExpr::Param("p")])),
        DimMatcher::Expr(DimExpr::Pow(&DimExpr::Param("z"), 3)),
        DimMatcher::Expr(DimExpr::Param("c")),
    ]);

    let batch = 2;
    let height = 3;
    let width = 2;
    let padding = 4;
    let channels = 5;
    let z = 4;

    let shape = [
        12,
        batch,
        1,
        2,
        3,
        height * padding,
        width * padding,
        z * z * z,
        channels,
    ];
    let env = [("p", padding), ("c", channels)];

    b.iter(|| {
        PATTERN.assert_shape_every_n(&shape, &env, 10);
    });
}