fastvideo.v1.layers.custom_op

fastvideo.v1.layers.custom_op#

Module Contents#

Classes#

CustomOp

Base class for custom ops. Dispatches the forward method to the appropriate backend.

Data#

logger

API#

class fastvideo.v1.layers.custom_op.CustomOp[source]#

Bases: torch.nn.Module

Base class for custom ops. Dispatches the forward method to the appropriate backend.

Initialization

Initialize internal Module state, shared by both nn.Module and ScriptModule.

abstract static default_on() → bool[source]#

On by default if level < CompilationLevel.PIECEWISE Specifying ‘all’ or ‘none’ in custom_op takes precedence.

dispatch_forward() → collections.abc.Callable[source]#
classmethod enabled() → bool[source]#
forward(*args, **kwargs) → Any[source]#
forward_cpu(*args, **kwargs) → Any[source]#
abstract forward_cuda(*args, **kwargs) → Any[source]#
abstract forward_native(*args, **kwargs) → Any[source]#

PyTorch-native implementation of the forward method. This method is optional. If implemented, it can be used with compilers such as torch.compile or PyTorch XLA. Also, it can be used for testing purposes.

forward_oot(*args, **kwargs) → Any[source]#
forward_tpu(*args, **kwargs) → Any[source]#
op_registry: dict[str, type[fastvideo.v1.layers.custom_op.CustomOp]][source]#

None

classmethod register(name: str) → collections.abc.Callable[source]#
fastvideo.v1.layers.custom_op.logger[source]#

‘init_logger(…)’

Contents