Lorax.Shape (lorax v0.2.1)

Summary

Functions

Infers the shape of the LoRA matrices. Only supports :dense and :conv operations

Functions

Link to this function

calc_ab(op, r, parameters)

Infers the shape of the LoRA matrices. Only supports :dense and :conv operations

Dense LoRA kernels

Suppose we have a target node W and input tensor x, During injection, the A matrix will project the input tensor down to r-dimensional space. Afterwards, the B matrix will project the result back up to some unknown dimensionality. To figure out the output dimensionality, we inspect W's kernel shape.

Convolution LoRA kernels

In addition to figuring out the output dimensionality, we need to retrieve the convolution kernel_size of W. When calling Axon.conv, it's passed as an option, but then is no longer stored inside the Axon node. To figure out the kernel size, we also inspect W's kernel shape.