easydel.infra.etils

EasyDeL utilities for enumerations and type definitions.

This module provides enumerations and type definitions used throughout the EasyDeL framework for configuration and runtime options. It includes enums for optimizers, schedulers, quantization methods, platforms, backends, and various configuration options.

Enumerations:

EasyDeLOptimizers: Available optimization algorithms EasyDeLSchedulers: Learning rate scheduling strategies EasyDeLGradientCheckPointers: Gradient checkpointing methods EasyDeLPlatforms: Kernel execution platforms EasyDeLBackends: JAX backend hardware targets

Type Aliases:

AVAILABLE_GRADIENT_CHECKPOINTS: Literal type for gradient checkpoint options AVAILABLE_SCHEDULERS: Literal type for scheduler options AVAILABLE_OPTIMIZERS: Literal type for optimizer options AVAILABLE_ATTENTION_MECHANISMS: Literal type for attention implementations AVAILABLE_SPARSE_MODULE_TYPES: Literal type for sparse matrix formats AVAILABLE_GRADIENT_CHECKPOINT_TARGETS: Literal type for checkpoint names used in models

Constants:

DEFAULT_ATTENTION_MECHANISM: Default attention mechanism to use

Functions:

define_flags_with_default: Create argparse flags from default values

Example

>>> from easydel.infra.etils import EasyDeLOptimizers, EasyDeLBackends
>>>
>>> # Use enums for configuration
>>> optimizer = EasyDeLOptimizers.ADAMW
>>> backend = EasyDeLBackends.TPU
>>>
>>> # Parse command-line arguments with defaults
>>> args, defaults = define_flags_with_default(
...     learning_rate=1e-3,
...     batch_size=32,
...     _required_fields=["model_name"]
... )

class easydel.infra.etils.EasyDeLBackends(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration of JAX backend types supported by EasyDeL.

Specifies the target hardware device type for JAX computations.

CPU

Use the CPU backend.

GPU

Use the GPU backend.

TPU

Use the TPU backend.

TT

Use the Tenstorrent backend.

CPU = 'cpu'

GPU = 'gpu'

TPU = 'tpu'

TT = 'tt'

class easydel.infra.etils.EasyDeLGradientCheckPointers(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration of gradient checkpointing strategies available in EasyDeL.

Gradient checkpointing is a technique to reduce memory usage during training by recomputing activations during the backward pass instead of storing them.

EVERYTHING_SAVEABLE

Checkpoints residuals, attentions, and hidden states. This is the most memory-intensive checkpointing strategy.

NOTHING_SAVEABLE

Checkpoints only the residuals. This strategy saves the most memory but requires more recomputation.

CHECKPOINT_DOTS

Checkpoints matrix multiplications and intermediate activations.

CHECKPOINT_DOTS_WITH_NO_BATCH_DMIS

Similar to CHECKPOINT_DOTS but avoids checkpointing operations involving batch dimensions.

NONE

No gradient checkpointing is applied.

CHECKPOINT_DOTS = 'checkpoint_dots'

CHECKPOINT_DOTS_WITH_NO_BATCH_DMIS = 'checkpoint_dots_with_no_batch_dims'

DOTS_SAVEABLE = 'dots_saveable'

DOTS_WITH_NO_BATCH_DIMS_AVAILABLE = 'dots_with_no_batch_dims_saveable'

EVERYTHING_SAVEABLE = 'everything_saveable'

NONE = ''

NOTHING_SAVEABLE = 'nothing_saveable'

SAVE_ANYTHING_EXCEPT_THESE_NAMES = 'save_anything_except_these_names'

SAVE_ANY_NAMES_BUT_THESE = 'save_any_names_but_these'

SAVE_FROM_BOTH_POLICIES = 'save_from_both_policies'

SAVE_ONLY_THESE_NAMES = 'save_only_these_names'

class easydel.infra.etils.EasyDeLOptimizers(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration of available optimizers in the EasyDeL library.

ADAFACTOR

Represents the Adafactor optimizer.

LION

Represents the Lion optimizer.

ADAMW

Represents the AdamW optimizer.

RMSPROP

Represents the RMSprop optimizer.

ADAFACTOR = 'adafactor'

ADAMW = 'adamw'

LION = 'lion'

MARS = 'mars'

MUON = 'muon'

QUAD = 'quad'

RMSPROP = 'rmsprop'

SKEW = 'skew'

class easydel.infra.etils.EasyDeLPlatforms(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration of platforms or kernel execution backends supported by EasyDeL.

This allows selecting optimized kernel implementations for different hardware or software environments.

JAX

Use standard JAX kernel implementations.

TRITON

Use Triton-based kernel implementations (often for GPUs).

PALLAS

Use Pallas-based kernel implementations (often for TPUs).

JAX = 'jax'

PALLAS = 'pallas'

TRITON = 'triton'

class easydel.infra.etils.EasyDeLSchedulers(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration of available learning rate schedulers in EasyDeL.

NONE

Indicates no scheduler should be used.

LINEAR

Represents a linear learning rate decay scheduler.

COSINE

Represents a cosine annealing learning rate scheduler.

COSINE = 'cosine'

LINEAR = 'linear'

NONE = 'None'

class easydel.infra.etils.StoreTupleAction(option_strings, dest, nargs=None, const=None, default=None, type=None, choices=None, required=False, help=None, metavar=None)

Bases: Action

Custom argparse action to parse a comma-separated string into a tuple of integers.

This action is used by define_flags_with_default when a default value is a tuple. It takes the comma-separated string provided on the command line and attempts to convert each part into an integer, storing the result as a tuple in the namespace.

Raises

argparse.ArgumentTypeError – If the provided value cannot be parsed as a comma-separated list of integers.

easydel.infra.etils.define_flags_with_default(_required_fields: list | None = None, **kwargs) → tuple[argparse.Namespace, dict[str, Any]]

Defines command-line flags using argparse based on provided keyword arguments.

This function dynamically creates argparse arguments for each key-value pair in kwargs. It infers the argument type from the default value and handles tuple types specifically. It also supports marking certain fields as required.

Parameters

• _required_fields (tp.List, optional) – A list of flag names that are mandatory. An error will be raised if these flags are not provided or are empty strings. Defaults to None.

• **kwargs – Keyword arguments where keys are flag names (without –) and values are their default values.

Returns

A tuple containing:

• An argparse.Namespace object holding the parsed command-line arguments.

• A dictionary mapping the original flag names to their default values.

Return type

tp.Tuple[argparse.Namespace, tp.Dict[str, tp.Any]]

Raises

ValueError – If a required field (from _required_fields) is not provided or is an empty string on the command line.