module

ML::ShapeOps

Utility functions for shape operations

Instance methods

batched_matmul_shape(a : Shape, b : Shape) : Shape

Compute output shape for batched matmul: [..., M, K] @ [..., K, N] -> [..., M, N]

Source
flatten_shape(shape : Shape, start_dim : Int32 = 0, end_dim : Int32 = -1) : Shape

Flatten shape from start_dim to end_dim

Source
matmul_shape(a : Shape, b : Shape) : Shape

Compute output shape for matmul: [M, K] @ [K, N] -> [M, N]

Source
squeeze_shape(shape : Shape, dim : Int32 | Nil = nil) : Shape

Squeeze: remove dimensions of size 1

Source
transpose_shape(shape : Shape) : Shape

Transpose last two dimensions

Source
unsqueeze_shape(shape : Shape, dim : Int32) : Shape

Unsqueeze: add dimension of size 1 at position

Source