Torch Extensions

Piecewise Linear Function

class teachcompute.torch_extensions.piecewise_linear.PiecewiseLinearFunction(*args, **kwargs)

Implements a function similar to a piecewise linear function. It multiplies by different coefficient on negative and positive number. It takes a tensor of shape (N, 1).

class teachcompute.torch_extensions.piecewise_linear.PiecewiseLinearFunctionC(*args, **kwargs)

Same function as PiecewiseLinearFunction but the implementation of forward and backward functions are done in C. See piecewise_linear_forward and piecewise_linear_backward. It follows the tutorial Custom C++ and CUDA Extensions.

teachcompute.torch_extensions.piecewise_linear_c.piecewise_linear_backward(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor) → list[torch.Tensor]

PiecewiseLinearC backward

teachcompute.torch_extensions.piecewise_linear_c.piecewise_linear_forward(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor) → list[torch.Tensor]

PiecewiseLinearC forward