context

Defines the context for Tricycle operations.

This module provides a dataclass for storing Tricycle context information, including mixed precision and loss scaling settings.

class TricycleContext(use_mixed_precision=False, loss_scale_factor=128)

Bases: object

A dataclass to store Tricycle context information.

Parameters:

• use_mixed_precision (bool)

• loss_scale_factor (int)

use_mixed_precision

Flag to enable mixed precision. Default is False. Note: It’s recommended to use the tricycle/utils.py:UseMixedPrecision context manager for mixed precision training instead of modifying this directly.

Type:

bool

loss_scale_factor

Factor to scale the loss when using mixed precision. This helps prevent under and overflowing. Default is 128.

Type:

int

loss_scale_factor: int = 128

use_mixed_precision: bool = False