easydel.infra.etils

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.

CPU = 'cpu'

GPU = 'gpu'

TPU = 'tpu'

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'

EVERYTHING_SAVEABLE = 'everything_saveable'

NONE = ''

NOTHING_SAVEABLE = 'nothing_saveable'

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'

RMSPROP = 'rmsprop'

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.EasyDeLQuantizationMethods(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration of quantization methods supported by EasyDeL.

Quantization reduces the precision of model weights and/or activations to save memory and potentially speed up inference.

NONE

No quantization is applied.

NF4

Represents NormalFloat 4-bit quantization.

A8BIT

Represents 8-bit affine quantization.

A8BIT = '8bit'

NF4 = 'nf4'

NONE = 'None'

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, **kwargs) → Tuple[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.