nytorch package

nytorch.base module

class nytorch.base.NytoModuleBase

Bases: Module

Base class for NytoModule.

Acts as an interface for ParticleDataImp and various updaters.

Attributes:

_module_id (ModuleID):

ID of the module.

_particle_kernel (Optional[ParticleKernelImp]):

Points to the particle kernel; set to None when wrapped by ParticleModule.

property is_root: bool

Check if the current module is the root module.

Returns:

bool: True if the module is the root module, False otherwise.

class nytorch.base.ParticleDataImp(modules: OrderedDict[int, Module], params: OrderedDict[int, Parameter])

Bases: ParticleData[VersionDataImp, ParticleDataImp]

Implementation of ParticleData.

This class stores references to all parameters of corresponding particles and is accessed when particles undergo particle operations. Instances are stored in ParticleKernel.data.

Args:

modules (ModuleDict):

Contains references to all modules within particles.

params (ParamDict):

Contains references to all parameters within particles.

Attributes:

modules (ModuleDict):

Contains references to all modules within particles.

params (ParamDict):

Contains references to all parameters within particles.

add_modules(modules: OrderedDict[int, Module]) → None

Add modules to the current ParticleDataImp instance.

The keys in the passed modules should not already exist. Offsets can be applied to all keys in modules if needed.

Args:

modules (ModuleDict): The modules to be added.

add_params(params: OrderedDict[int, Parameter]) → None

Add params to the current ParticleDataImp instance.

The keys in the passed params should not already exist. Offsets can be applied to all keys in params if needed.

Args:

params (ParamDict): The params to be added.

assign(params: OrderedDict[int, Parameter]) → None

Replace current params with the given params.

The keys in the passed params must match the current params.

Args:

params (ParamDict): The new params to be assigned.

copy(vdata: VersionDataImp) → ParticleDataImp

Creates a shallow copy of the current particle data instance.

This method returns a new instance that shares references to the same parameters as the current instance.

Args:

vdata (Tvdata): The corresponding VersionData instance.

Returns:

Tpdata: A shallow copy of the particle data instance.

create(vdata: VersionDataImp, params: OrderedDict[int, Parameter]) → ParticleDataImp

Create a new ParticleDataImp with the same structure but different params.

Args:

vdata (VersionDataImp):

Used to guide the establishment of references between modules and params.

params (ParamDict):

Params for the new particle.

Returns:

ParticleDataImp: The newly created ParticleDataImp instance.

init_kernel(kernel: ParticleKernel) → None

Initializes the particle kernel before creating a new particle.

This method ensures that the ParticleKernel is properly initialized and ready to store this instance.

Args:

kernel (ParticleKernel): The ParticleKernel instance that will store this instance.

modules: OrderedDict[int, Module]

params: OrderedDict[int, Parameter]

remove_modules(modules: set[ModuleID]) → None

Remove modules from the current ParticleDataImp instance.

Args:

modules (set[ModuleID]): The set of ModuleIDs to be removed.

remove_params(params: set[ParamID]) → None

Remove params from the current ParticleDataImp instance.

Args:

params (set[ParamID]): The set of ParamIDs to be removed.

class nytorch.base.VersionDataImp(meta: OrderedDict[int, ModuleMeta], config: OrderedDict[int, ParamConfig])

Bases: VersionData[VersionDataImp, ParticleDataImp]

Implementation of VersionData.

This class stores metadata of particles, or more precisely, metadata of species where particles belong. When particles undergo particle operations, this instance needs to be accessed. Instances are stored in VersionKernel.data.

The content of an instance is immutable once determined, and a copy will be made when new metadata for particles is needed, and modifications will be made on the new instance.

Args:

meta (MetaDict):

Stores references of modules to modules and parameters within modules.

config (ConfigDict):

Stores settings of particle operations for all parameters.

Attributes:

meta (MetaDict):

Stores references of modules to modules and parameters within modules.

config (ConfigDict):

Stores settings of particle operations for all parameters.

config: OrderedDict[int, ParamConfig]

copy() → VersionDataImp

Creates a deep copy of the instance.

This method returns a deep copy of the current instance, ensuring that all nested data structures are also copied.

Returns:

Tvdata: A deep copy of the instance.

get_garbage_modules() → set[ModuleID]

Retrieve module IDs that are unreachable from the root module.

Returns:

Set[ModuleID]: A set of module IDs that are considered garbage.

get_garbage_params() → set[ParamID]

Retrieve parameter IDs that are unreachable from the root module.

Returns:

Set[ParamID]: A set of parameter IDs that are considered garbage.

get_sub_modules(root: ModuleID) → set[ModuleID]

Retrieve all module IDs (including the root) under the target module.

Args:

root (ModuleID): The ModuleID of the target module.

Returns:

Set[ModuleID]: A set of all module IDs under the target module.

get_sub_params(root: ModuleID) → set[ParamID]

Retrieve all parameter IDs (including the root) under the target module.

Args:

root (ModuleID): The ModuleID of the target module.

Returns:

Set[ParamID]: A set of all parameter IDs under the target module.

init_kernel(kernel: VersionKernel) → None

Initializes the version kernel before creating a new version.

This method ensures that the VersionKernel is properly initialized and ready to store this instance.

Args:

kernel (VersionKernel): The VersionKernel instance that will store this instance.

init_particle(pdata: ParticleDataImp) → None

Ensures the keys of pdata are in order before the particle enters the new version.

This method verifies and organizes the keys of the particle data before it is added to the new version.

Args:

pdata (Tpdata): The ParticleData subclass instance of the particle entering the new version.

meta: OrderedDict[int, ModuleMeta]

remove_modules(modules: set[ModuleID]) → None

Remove metadata of specified modules.

Args:

modules (Set[ModuleID]): A set of specified module IDs to remove.

remove_params(params: set[ParamID]) → None

Remove metadata of specified parameters.

Args:

params (Set[ParamID]): A set of specified parameter IDs to remove.

nytorch.base.create_new_group(ref: Module) → ParticleKernel[VersionDataImp, ParticleDataImp]

Create a new version kernel and a new particle kernel.

Note

The passed ref Module must be in a state where metadata is empty, meaning it should not have any parameters or submodules.

Args:

ref (Module): The reference module to base the new group on.

Returns:

ParticleKernelImp: The created particle kernel.

nytorch.kernel module

class nytorch.kernel.Particle(*args, **kwds)

Bases: ABC, Generic[Tproduct]

Abstract base class to be used in conjunction with a corresponding Product subclass.

Subclasses of this class must implement transformation methods (product and product_), which automatically perform particle operations through the corresponding Product subclass. The corresponding Product subclass should also be defined to allow transformation between the two using the product and particle methods.

Example usage:

class ProductImp(Product[ParticleImp]):
    ...

class ParticleImp(Particle[ProductImp]):
    ...

particle: ParticleImp = ParticleImp()
product: ProductImp = particle.product()
new_particle: ParticleImp = product.particle()

clone() → Tparticle

Clone the current particle.

Returns:

Tparticle: A new particle that is a clone of the current particle.

abstract product() → Tproduct

Transforms the current particle into its corresponding Product instance.

Returns:

Tproduct: The corresponding Product instance.

abstract product_(product: Tproduct) → Self

Transforms the current particle into a new particle instance based on the given Product.

Args:

product (Tproduct): The corresponding Product instance containing parameters to import.

Returns:

Self: The current particle instance after parameter importation.

rand() → Tparticle

Generate a random particle with uniformly distributed values.

Returns:

Tparticle: A new particle with randomly generated values.

randn() → Tparticle

Generate a random particle with normally distributed values.

Returns:

Tparticle: A new particle with randomly generated values.

class nytorch.kernel.ParticleData(*args, **kwds)

Bases: ABC, Generic[Tvdata, Tpdata]

Abstract base class for storing parameters of particles.

This class stores references to all parameters of the corresponding particle. It is essential for particle operations and will be stored in ParticleKernel.data.

A corresponding subclass of VersionData should be defined when instantiating this class, as shown below:

class VersionDataImp(VersionData[VersionDataImp, ParticleDataImp]):
    ...

class ParticleDataImp(ParticleData[VersionDataImp, ParticleDataImp]):
    ...

abstract copy(vdata: Tvdata) → Tpdata

Creates a shallow copy of the current particle data instance.

This method returns a new instance that shares references to the same parameters as the current instance.

Args:

vdata (Tvdata): The corresponding VersionData instance.

Returns:

Tpdata: A shallow copy of the particle data instance.

abstract init_kernel(kernel: ParticleKernel) → None

Initializes the particle kernel before creating a new particle.

This method ensures that the ParticleKernel is properly initialized and ready to store this instance.

Args:

kernel (ParticleKernel): The ParticleKernel instance that will store this instance.

class nytorch.kernel.ParticleKernel(version: VersionKernel[Tvdata, Tpdata], data: Tpdata)

Bases: Generic[Tvdata, Tpdata]

Core class for managing particles.

Each ParticleKernel instance represents a particle, storing the particle’s specific parameters. The particle’s metadata is stored in the version instance pointed to by version.

Args:

version (VersionKernel[Tvdata, Tpdata]):

Points to the corresponding version.

data (Tpdata):

The ParticleData subclass instance storing particle parameters, used to assist particle operations.

Attributes:

version (VersionKernel[Tvdata, Tpdata]):

Points to the corresponding version.

data (Tpdata):

The ParticleData subclass instance storing particle parameters, used to assist particle operations.

create(data: Tpdata) → ParticleKernel[Tvdata, Tpdata]

Creates a ParticleKernel belonging to the same species.

Args:

data (Tpdata): The ParticleData subclass instance storing particle parameters.

Returns:

ParticleKernel[Tvdata, Tpdata]: A new ParticleKernel instance.

data: Tpdata

detach() → ParticleKernel[Tvdata, Tpdata]

Detaches the current particle kernel instance to a new species.

Creates a deep copy of the version kernel and a shallow copy of the current particle kernel, then associates the copied version kernel with the new particle kernel instance.

Returns:

ParticleKernel[Tvdata, Tpdata]: The detached ParticleKernel instance.

particle_update() → None

Updates the particle to the newest version.

This method ensures the particle is updated through all intermediate versions to the latest version.

version: VersionKernel[Tvdata, Tpdata]

version_update(version_updater: VersionUpdater[Tvdata, Tpdata]) → None

Creates a new version and updates the particle.

This method ensures that if updating the version succeeds but updating the particle fails, the version update is rolled back.

Args:

version_updater (VersionUpdater[Tvdata, Tpdata]): The VersionUpdater subclass instance used to update the version.

class nytorch.kernel.ParticleUpdater(*args, **kwds)

Bases: ABC, Generic[Tvdata, Tpdata]

Abstract base class for tools that update ParticleKernel instances.

This class provides a template for implementing tools that handle direct operations on particles when updating to the next version. The run method defines these operations.

The following methods are called in sequence when updating a ParticleKernel instance:

1. set_version_data

2. set_next_version_data

3. set_particle_data

4. run

Methods 1, 2, and 3 are called once during initialization, while method 4 is called once for each particle that needs to be updated to the new version.

Subclasses must implement these methods to define the update process.

abstract run(pdata: Tpdata) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

abstract set_next_version_data(vdata: Tvdata) → ParticleUpdater[Tvdata, Tpdata]

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

abstract set_particle_data(pdata: Tpdata) → ParticleUpdater[Tvdata, Tpdata]

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

abstract set_version_data(vdata: Tvdata) → ParticleUpdater[Tvdata, Tpdata]

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

class nytorch.kernel.Product(*args, **kwds)

Bases: ABC, Generic[Tparticle]

Abstract base class to be used in conjunction with a corresponding Particle subclass.

Subclasses of this class must implement particle operations and transformation methods. The corresponding Particle subclass should also be defined to allow transformation between the two using the product and particle methods.

Example usage:

class ProductImp(Product[ParticleImp]):
    ...

class ParticleImp(Particle[ProductImp]):
    ...

particle: ParticleImp = ParticleImp()
product: ProductImp = particle.product()
new_particle: ParticleImp = product.particle()

abstract clone() → Tproduct

Clone the current product.

Returns:

Tproduct: A new product that is a clone of the current product.

abstract particle() → Tparticle

Transforms the current particle into its corresponding Product instance.

Returns:

Tparticle: The corresponding Particle instance.

abstract rand() → Tproduct

Generate a random product with uniformly distributed values.

Returns:

Tproduct: A new product with randomly generated values.

abstract randn() → Tproduct

Generate a random product with normally distributed values.

Returns:

Tproduct: A new product with randomly generated values.

class nytorch.kernel.VersionData(*args, **kwds)

Bases: ABC, Generic[Tvdata, Tpdata]

Abstract base class for storing metadata of particles.

This class is responsible for storing the metadata of particles, specifically the metadata of the species to which the particle belongs. It is essential for particle operations and will be stored in VersionKernel.data.

Once an instance is created, its content remains unchanged. When new metadata is required, a copy is created and modifications are applied to the new instance.

A corresponding subclass of ParticleData should be defined when instantiating this class, as shown below:

class VersionDataImp(VersionData[VersionDataImp, ParticleDataImp]):
    ...

class ParticleDataImp(ParticleData[VersionDataImp, ParticleDataImp]):
    ...

abstract copy() → Tvdata

Creates a deep copy of the instance.

This method returns a deep copy of the current instance, ensuring that all nested data structures are also copied.

Returns:

Tvdata: A deep copy of the instance.

abstract init_kernel(kernel: VersionKernel) → None

Initializes the version kernel before creating a new version.

This method ensures that the VersionKernel is properly initialized and ready to store this instance.

Args:

kernel (VersionKernel): The VersionKernel instance that will store this instance.

abstract init_particle(pdata: Tpdata) → None

Ensures the keys of pdata are in order before the particle enters the new version.

This method verifies and organizes the keys of the particle data before it is added to the new version.

Args:

pdata (Tpdata): The ParticleData subclass instance of the particle entering the new version.

class nytorch.kernel.VersionKernel(data: Tvdata)

Bases: Generic[Tvdata, Tpdata]

Core class for managing versions of particle metadata.

Each VersionKernel instance represents a version or state of particle metadata. When the metadata state of a particle changes, a new VersionKernel instance is created and linked via next_version. The particle_updater facilitates updates to particles within the current version.

Args:

data (Tvdata):

The VersionData subclass instance storing particle metadata, used to assist particle operations.

Attributes:

next_version (Optional[VersionKernel[Tvdata, Tpdata]]):

Points to the next version.

particle_updater (Optional[ParticleUpdater[Tvdata, Tpdata]]):

Tool for updating particles within the current version.

data (Tvdata):

The VersionData subclass instance storing particle metadata, used to assist particle operations.

copy() → VersionKernel

Creates a shallow copy of the current version.

Returns:

VersionKernel: A shallow copy of the current VersionKernel.

data: Tvdata

del_next_version() → None

Removes the new version.

This method is called after updating the version if an error occurs when updating particles, indicating that there is a problem with the new version, and the new version will be removed.

property is_newest: bool

Checks if the current version is the newest version.

Returns:

bool: True if the current version is the newest, False otherwise.

next_version: Optional[VersionKernel[Tvdata, Tpdata]]

particle_update(particle_data: Tpdata) → None

Updates the particle data to the new version using the particle updater.

Args:

particle_data (Tpdata): The instance of ParticleData subclass to be updated.

Raises:

AssertionError: If particle_updater or next_version is None.

particle_updater: Optional[ParticleUpdater[Tvdata, Tpdata]]

version_update(version_updater: VersionUpdater[Tvdata, Tpdata], particle_data: Tpdata) → None

Generates a new version.

Args:

version_updater (VersionUpdater[Tvdata, Tpdata]):

The VersionUpdater subclass instance that generates next_version and particle_updater.

particle_data (Tpdata):

The ParticleData subclass instance of the particle initiating this operation.

Raises:

AssertionError: If particle_updater or next_version is None.

class nytorch.kernel.VersionUpdater(*args, **kwds)

Bases: ABC, Generic[Tvdata, Tpdata]

Abstract base class for tools that update VersionKernel instances.

This class provides a template for implementing tools that handle updates of VersionKernel instances. A corresponding subclass of ParticleUpdater should be implemented to manage the particle updates.

The following methods are called in sequence when updating a VersionKernel instance:

1. set_version_data

2. set_particle_data

3. run

Subclasses must implement these methods to define the update process.

abstract run(vdata: Tvdata) → tuple[Tvdata, 'ParticleUpdater']

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

abstract set_particle_data(pdata: Tpdata) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

abstract set_version_data(vdata: Tvdata) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

nytorch.module module

class nytorch.module.NytoModule

Bases: NytoModuleBase, Particle[ParamProduct]

Implementation of NytoModuleBase, compatible with torch.nn.Module functionality.

Features:

• Particle operations.

• Transformation to ParamProduct.

• Version management.

Example1: Particle Operations:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_param = nn.Parameter(torch.Tensor([0., 1., 2.]))

net = Net()
new_net = 2 * net

>>> new_net.my_param
Parameter containing:
tensor([0., 2., 4.], requires_grad=True)

Example2: Transform to ParamProduct:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_param = nn.Parameter(torch.Tensor([0., 1., 2.]))

net = Net()
product = net.product()
assert isinstance(product, ParamProduct)

Example3: Version Management:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_param = nn.Parameter(torch.Tensor([0., 1., 2.]))

net = Net()
new_net = 2 * net
new_net.my_param2 = nn.Parameter(torch.Tensor([3., 4.]))
net.touch()

>>> net.my_param2
Parameter containing:
tensor([3., 4.], requires_grad=True)

add_module(name: str, module: Optional[Module]) → None

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (string): name of the child module. The child module can be

accessed from this module using the given name

module (Module): child module to be added to the module.

apply_param_config(fn: Callable[[int, ParamConfig], None], name: Optional[str] = None) → None

Manually modify the ParamConfig instance.

Args:

fn (Callable[[ParamID, ParamConfig], None]):

Function that takes a param ID and its corresponding ParamConfig instance, and modifies the configuration.

name (str):

Name of the attribute to modify the configuration for, default is to modify the configuration for all attributes.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_module = torch.nn.Linear(3, 2)
        self.my_param = torch.nn.Parameter(torch.randn(2, 2))

def set_operational_false(pid: ParamID, conf: ParamConfig):
    conf.operational = False

net = Net()
net.apply_param_config(set_operational_false, "my_module")

clone_from(source: Module) → Tmodule

Clone another particle from a different species into the current species.

Args:

attr_name (str): Particle from another species.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_param = torch.nn.Parameter(torch.randn(2, 2))

net1 = Net()
net2 = Net()
net3 = net1.clone_from(net2)

>>> net1._version_kernel is net2._version_kernel
False

>>> net1._version_kernel is net3._version_kernel
True

>>> net1.my_param is net3.my_param
False

detach() → Tmodule

Copy the parameters of the current particle to a new species.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_param = torch.nn.Parameter(torch.randn(2, 2))

net1 = Net()
net2 = net1.detach()

>>> net1._version_kernel is net2._version_kernel
False 

>>> net1.my_param is net2.my_param
True

get_param_id(target_param: Parameter) → int

Find the ID of the specified parameter.

Args:

target_param (ParamType):

The specified parameter.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_module = torch.nn.Linear(3, 2)
        self.my_param = torch.nn.Parameter(torch.randn(2, 2))

net = Net()

>>> net.get_param_id(net.my_param)
2

product() → ParamProduct

Transform into a ParamProduct to optimize particle operations.

Returns:

ParamProduct: The ParamProduct instance representing the module.

product_(product: ParamProduct) → Self

Import the parameters of a ParamProduct into the current module.

Args:

product (ParamProduct): The ParamProduct instance containing parameters to import.

Returns:

Self: The current module after importing the parameters.

register_buffer(name: str, tensor: Optional[Tensor], persistent: bool = True) → None

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (string): name of the buffer. The buffer can be accessed

from this module using the given name

tensor (Tensor): buffer to be registered. persistent (bool): whether the buffer is part of this module’s

state_dict.

Example:

>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_parameter(name: str, param: Optional[Parameter]) → None

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (string): name of the parameter. The parameter can be accessed

from this module using the given name

param (Parameter): parameter to be added to the module.

set_param_config(operational: Optional[bool] = None, clone: Optional[bool] = None, name: Optional[str] = None) → None

Modify the default configuration of the ParamConfig instance.

Args:

operational (Optional[bool]):

If True, the parameter is included in particle operations.

clone (Optional[bool]):

If False, the parameter is cloned if not included in particle operations.

name (Optional[str]):

Name of the attribute to modify the configuration for, default is to modify the configuration for all attributes.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_module = torch.nn.Linear(3, 2)
        self.my_param = torch.nn.Parameter(torch.randn(2, 2))

net = Net()
net.set_param_config(operational=False, name="my_module")

touch() → Self

Update the particle to the latest version.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()

net1 = Net()
net2 = net1.clone()
net1.my_param = torch.nn.Parameter(torch.randn(2, 2))

>>> hasattr(net2, "my_param")
False

>>> net2.touch()
>>> hasattr(net2, "my_param")
True

class nytorch.module.ParamProduct(kernel: ParticleKernel[VersionDataImp, ParticleDataImp], module_id: int, params: OrderedDict[int, Parameter])

Bases: Product[Tmodule]

Implementation of Product for NytoModule instances.

Implements particle operations and transformation to NytoModule.

Args:

kernel (ParticleKernelImp):

Reference to the particle kernel.

module_id (ModuleID):

Specifies which module to return when transforming to NytoModule.

params (ParamDict):

Model parameters for particle operations.

Attributes:

kernel (ParticleKernelImp):

Reference to the particle kernel.

module_id (ModuleID):

Specifies which module to return when transforming to NytoModule.

params (ParamDict):

Model parameters for particle operations.

binary_operator(other: ParamProduct, fn: Callable[[Parameter, Parameter, ParamConfig], Parameter]) → ParamProduct

Perform a binary operation between two ParamProduct instances.

Args:

other (ParamProduct):

Another ParamProduct instance involved in the binary operation.

fn (Callable[[ParamType, ParamType, ParamConfig], ParamType]):

Binary operation logic.

Returns:

ParamProduct: New ParamProduct instance after the binary operation.

Raises:

AssertionError: If other.kernel.version is not self.kernel.version.

Note

In writing the function fn, gradient calculation does not need to be disabled because torch.no_grad is used within the unary_operator() method to disable gradient calculation.

Note

The source of other must belong to the same species as self, which can be checked as follows:

assert other.kernel.version is self.kernel.version

Example:

class Net(NytoModule):
    def __init__(self, my_tensor):
        super().__init__()
        self.my_param = nn.Parameter(my_tensor)

net1 = Net(torch.Tensor([0., 1., 2.]))
net2 = net1.clone_from(Net(torch.Tensor([5., 4., 3.])))
product1 = net1.product()
product2 = net2.product()

fn = lambda param1, param2, conf: nn.Parameter(param1-param2)
new_product = product1.binary_operator(product2, fn)
new_net = new_product.module()

>>> new_net.my_param
Parameter containing:
tensor([-5., -3., -1.], requires_grad=True)

clone() → ParamProduct

Clone the current product.

Returns:

Tproduct: A new product that is a clone of the current product.

kernel: ParticleKernel[VersionDataImp, ParticleDataImp]

module() → Tmodule

Transforms into the corresponding NytoModule subclass.

module_id: int

params: OrderedDict[int, Parameter]

particle() → Tmodule

Transforms into the corresponding NytoModule subclass.

property pdata: ParticleDataImp

Parameter data of the particle.

rand() → ParamProduct

Generate a random product with uniformly distributed values.

Returns:

Tproduct: A new product with randomly generated values.

randn() → ParamProduct

Generate a random product with normally distributed values.

Returns:

Tproduct: A new product with randomly generated values.

unary_operator(fn: Callable[[Parameter, ParamConfig], Parameter]) → ParamProduct

Perform a unary operation on the parameters.

Args:

fn (Callable[[ParamType, ParamConfig], ParamType]): Unary operation logic.

Returns:

ParamProduct: New ParamProduct instance after the unary operation.

Note

In writing the function fn, gradient calculation does not need to be disabled because torch.no_grad is used within the unary_operator() method to disable gradient calculation.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_param = nn.Parameter(torch.Tensor([0., 1., 2.]))

net = Net()
product = net.product()
new_product = product.unary_operator(lambda param, conf: nn.Parameter(param+10.))
new_net = new_product.module()

>>> new_net.my_param
Parameter containing:
tensor([10., 11., 12.], requires_grad=True)

property vdata: VersionDataImp

Version data of the particle.

nytorch.module.binary_lambda_wrapper(fn: Callable[[Parameter, Parameter], Tensor]) → Callable[[Parameter, Parameter, ParamConfig], Parameter]

Wraps a binary function for parameter operations.

Args:

fn (Callable[[ParamType, ParamType], torch.Tensor]): Binary operation function.

Returns:

Callable[[ParamType, ParamType, ParamConfig], ParamType]: Wrapped binary function.

nytorch.module.module_to_product(module: Tmodule) → ParamProduct

Transform a NytoModule instance into a ParamProduct instance.

Args:

module (Tmodule): The NytoModule instance to transform.

Returns:

ParamProduct: The corresponding ParamProduct instance.

Raises:

AssertionError: If module._particle_kernel is None.

nytorch.module.product_assign_to_module(product: ParamProduct, module: Tmodule) → None

Assign parameters from a ParamProduct instance to a NytoModule instance.

Args:

product (ParamProduct):

The ParamProduct instance containing parameters to assign.

module (Tmodule):

The NytoModule instance to assign parameters to.

Raises:

AssertionError: If product.kernel is not module._version_kernel or if module._particle_kernel is None.

nytorch.module.product_to_module(product: ParamProduct) → Module

Transform a ParamProduct instance back into a NytoModule instance.

Args:

product (ParamProduct): The ParamProduct instance to transform.

Returns:

Module: The corresponding NytoModule instance.

nytorch.module.unary_lambda_wrapper(fn: Callable[[Parameter], Tensor]) → Callable[[Parameter, ParamConfig], Parameter]

Wraps a unary function for parameter operations.

Args:

fn (Callable[[ParamType], torch.Tensor]): Unary operation function.

Returns:

Callable[[ParamType, ParamConfig], ParamType]: Wrapped unary function.

nytorch.mtype module

class nytorch.mtype.ModuleMeta(sub_modules: OrderedDict[str, Optional[ModuleID]], sub_params: OrderedDict[str, Optional[ParamID]])

Bases: object

Maintains mappings of module and parameter attribute names to their respective IDs.

Attributes:

sub_modules (OrderedDict[str, Optional[ModuleID]]):

Mapping of attribute names to submodule IDs.

sub_params (OrderedDict[str, Optional[ParamID]]):

Mapping of attribute names to parameter IDs.

copy() → ModuleMeta

Creates a deep copy of the current ModuleMeta instance.

Returns:

ModuleMeta: A deep copy of this ModuleMeta.

sub_modules: OrderedDict[str, Optional[ModuleID]]

sub_params: OrderedDict[str, Optional[ParamID]]

class nytorch.mtype.ParamConfig(operational: bool = True, clone: bool = True)

Bases: object

Stores configuration settings for parameter operations in particles.

By modifying these settings, users can influence particle behavior. Additional attributes can be added to support custom particle operations, which will define how these custom attributes are used.

Attributes:

operational (bool):

Indicates if the parameter is included in particle operations.

clone (bool):

Indicates if the parameter should be cloned when not included in particle operations.

clone: bool = True

copy() → ParamConfig

Creates a deep copy of the current ParamConfig instance.

Returns:

ParamConfig: A deep copy of this ParamConfig.

operational: bool = True

nytorch.particle_module module

class nytorch.particle_module.GetNytoModule(particle_module: ParticleModule[Tmodule])

Bases: Generic[Tmodule]

Context manager for automatic restoration and cleanup of kernel reference.

Args:

particle_module(ParticleModule[Tmodule]): The ParticleModule to be processed.

Attributes:

_module(ParticleModule[Tmodule]): The ParticleModule being processed.

To directly assign attributes and synchronization(touch) to an instance of a NytoModule subclass wrapped by ParticleModule, call restore_kernel_ref first to restore the kernel reference, and then call clear_kernel_ref after the assignment operation:

model = ParticleModule(MyNytoModule())
model_clone = model.clone()

# assign attributes
model.restore_kernel_ref()
model.root_module.data_embed = nn.Embedding(30, 8)
model.clear_kernel_ref()

# synchronization
model_clone.restore_kernel_ref()
model_clone.root_module.touch()
model_clone.clear_kernel_ref()

Alternatively, use GetNytoModule to simplify:

model = ParticleModule(MyNytoModule())
model_clone = model.clone()

with GetNytoModule(model) as my_nyto_module:
    my_nyto_module.data_embed = nn.Embedding(30, 8)

with GetNytoModule(model_clone) as my_nyto_module_clone:
    my_nyto_module_clone.touch()

class nytorch.particle_module.PMProduct(kernel: ParticleKernel[VersionDataImp, ParticleDataImp], module_id: int, params: OrderedDict[int, Parameter])

Bases: Product[ParticleModule]

Decorator for ParamProduct.

Implements particle operations and transforms into ParticleModule.

Args:

kernel (ParticleKernelImp):

Particle kernel instance.

module_id (ModuleID):

ID of the module.

params (ParamDict):

Parameters.

Attributes:

product (ParamProduct): Instance of ParamProduct.

binary_operator(other: PMProduct, fn: Callable[[Parameter, Parameter, ParamConfig], Parameter]) → PMProduct

Custom binary operation.

Args:

other (PMProduct):

Another ParamProduct instance participating in the binary operation.

fn (Callable[[ParamType, ParamType, ParamConfig], ParamType]):

Binary operation logic.

Returns:

PMProduct: Resultant PMProduct instance after applying binary operation.

Note

In writing the function fn, gradient calculation does not need to be disabled because torch.no_grad is used within the unary_operator() method to disable gradient calculation.

Note

The source of other must belong to the same species as the source of self, which can be checked as follows:

assert other.product.kernel.version is self.product.kernel.version

Example:

class Net(NytoModule):
    def __init__(self, my_tensor):
        super().__init__()
        self.my_param = nn.Parameter(my_tensor)

net1 = ParticleModule(Net(torch.Tensor([0., 1., 2.])))
net2 = ParticleModule(Net(torch.Tensor([5., 4., 3.])))
net2 = net1.clone_from(net2)
product1 = net1.product()
product2 = net2.product()

fn = lambda param1, param2, conf: nn.Parameter(param1-param2)
new_product = product1.binary_operator(product2, fn)
new_net = new_product.module()

>>> new_net.root_module.my_param
Parameter containing:
tensor([-5., -3., -1.], requires_grad=True)

clone() → PMProduct

Clone the current product.

Returns:

Tproduct: A new product that is a clone of the current product.

classmethod from_ParamProduct(product: ParamProduct) → PMProduct

Wrap a ParamProduct instance into a PMProduct instance.

Args:

product (ParamProduct): Wrapped ParamProduct instance.

Returns:

PMProduct: Wrapped PMProduct instance.

module() → ParticleModule

Transform into ParticleModule.

particle() → ParticleModule

Transform into ParticleModule.

product: ParamProduct

rand() → PMProduct

Generate a random product with uniformly distributed values.

Returns:

Tproduct: A new product with randomly generated values.

randn() → PMProduct

Generate a random product with normally distributed values.

Returns:

Tproduct: A new product with randomly generated values.

unary_operator(fn: Callable[[Parameter, ParamConfig], Parameter]) → PMProduct

Custom unary operation.

Args:

fn (Callable[[ParamType, ParamConfig], ParamType]): Unary operation logic.

Returns:

PMProduct: Resultant PMProduct instance after applying unary operation.

Note

In writing the function fn, gradient calculation does not need to be disabled because torch.no_grad is used within the unary_operator() method to disable gradient calculation.

Example:

class Net(NytoModule):
    def __init__(self):
        super().__init__()
        self.my_param = nn.Parameter(torch.Tensor([0., 1., 2.]))

net = ParticleModule(Net())
product = net.product()
new_product = product.unary_operator(lambda param, conf: nn.Parameter(param+10.))
new_net = new_product.module()

>>> new_net.root_module.my_param
Parameter containing:
tensor([10., 11., 12.], requires_grad=True)

class nytorch.particle_module.ParticleModule(root_module: Tmodule)

Bases: Module, Particle[PMProduct], Generic[Tmodule]

Decorator for NytoModule.

This class wraps a NytoModule to handle particle operations and transformations, addressing the issue of circular references by allowing the clearing and restoring of references to the particle kernel.

Features:

• Implements particle operations and transforms to PMProduct.

• Facilitates clearing and restoring the module’s reference to the particle kernel.

Note

Clearing the module’s reference to the particle kernel can eliminate circular references, reducing memory pressure.

Args:

root_module (Tmodule): The NytoModule instance to be wrapped.

Attributes:

particle_kernel (ParticleKernelImp):

Reference to the particle kernel for restoring the module’s reference.

root_module (Tmodule):

The root NytoModule being wrapped.

clear_kernel_ref() → None

Clears the module’s reference to the particle kernel.

This method is used to eliminate circular references and reduce memory usage.

clone_from(source: ParticleModule) → ParticleModule

Clone the particle from another ParticleModule instance.

Args:

source (ParticleModule): The ParticleModule instance to clone from.

Returns:

ParticleModule: The cloned ParticleModule instance.

forward(*args, **kwargs)

Forward pass through the root module.

Args:

*args:

Positional arguments for the root module’s forward method.

**kwargs:

Keyword arguments for the root module’s forward method.

Returns:

Tensor: The result of the root module’s forward method.

particle_kernel: ParticleKernelImp

product() → PMProduct

Transform the module into a PMProduct.

Returns:

PMProduct: A PMProduct instance representing the module.

product_(product: PMProduct) → Self

Import parameters from a PMProduct instance.

Args:

product (PMProduct): The PMProduct instance to import from.

Returns:

ParticleModule: The ParticleModule instance with imported parameters.

restore_kernel_ref() → None

Restores the module’s reference to the particle kernel.

This method is used to re-establish references that were cleared to reduce memory usage.

root_module: Tmodule

nytorch.particle_updater module

class nytorch.particle_updater.AddModuleParticleUpdater(module_id: ModuleID, attr_name: str, add_module_id: ModuleID, add_modules: ModuleDict, add_params: ParamDict, remove_modules: set[ModuleID], remove_params: set[ParamID])

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

add_module_id: ModuleID

add_modules: ModuleDict

add_params: ParamDict

attr_name: str

module_id: ModuleID

next_version_modules_meta: Optional[MetaDict]

owner: Optional[ParticleData]

remove_modules: set[ModuleID]

remove_params: set[ParamID]

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

class nytorch.particle_updater.AddParamParticleUpdater(module_id: ModuleID, attr_name: str, add_param_id: ParamID, remove_modules: set[ModuleID], remove_params: set[ParamID], add_param: Optional[ParamType] = None)

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

add_param: Optional[ParamType]

add_param_id: ParamID

attr_name: str

module_id: ModuleID

owner: Optional[ParticleDataImp]

remove_modules: set[ModuleID]

remove_params: set[ParamID]

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

class nytorch.particle_updater.DelBufferParticleUpdater(module_id: int, attr_name: str)

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

attr_name: str

module_id: int

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

class nytorch.particle_updater.DelModuleParticleUpdater(module_id: ModuleID, attr_name: str, remove_modules: set[ModuleID], remove_params: set[ParamID])

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

attr_name: str

module_id: ModuleID

remove_modules: set[ModuleID]

remove_params: set[ParamID]

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

class nytorch.particle_updater.DelParamParticleUpdater(module_id: ModuleID, attr_name: str, remove_params: set[ParamID])

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

attr_name: str

module_id: ModuleID

remove_params: set[ParamID]

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

class nytorch.particle_updater.RegisterBufferParticleUpdater(module_id: ModuleID, attr_name: str, value: Optional[torch.Tensor], persistent: bool, remove_modules: set[ModuleID], remove_params: set[ParamID])

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

attr_name: str

module_id: ModuleID

persistent: bool

remove_modules: set[ModuleID]

remove_params: set[ParamID]

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

value: Optional[torch.Tensor]

class nytorch.particle_updater.SetModuleNoneParticleUpdater(module_id: ModuleID, attr_name: str, remove_modules: set[ModuleID], remove_params: set[ParamID])

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

attr_name: str

module_id: ModuleID

remove_modules: set[ModuleID]

remove_params: set[ParamID]

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

class nytorch.particle_updater.SetParamNoneParticleUpdater(module_id: ModuleID, attr_name: str, remove_modules: set[ModuleID], remove_params: set[ParamID])

Bases: ParticleUpdater[VersionDataImp, ParticleDataImp]

attr_name: str

module_id: ModuleID

remove_modules: set[ModuleID]

remove_params: set[ParamID]

run(pdata: ParticleDataImp) → None

Executes the particle update process.

This method retrieves the data of the particle to be updated and performs the necessary operations to update the particle to the new version.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

set_next_version_data(vdata: VersionDataImp) → Self

Sets the next version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the next version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated next version data.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before performing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated particle data.

set_version_data(vdata: VersionDataImp) → Self

Sets the current version data required before performing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

ParticleUpdater[Tvdata, Tpdata]: The instance of the class with updated version data.

nytorch.utils module

nytorch.utils.clone_param(param: Parameter) → Parameter

Create a deep copy of the parameter.

Args:

param (ParamType): The parameter to be copied.

Returns:

ParamType: The deep copy of the parameter.

nytorch.utils.clone_params(params: OrderedDict[int, Parameter]) → OrderedDict[int, Parameter]

Create deep copies of parameters.

Args:

params (ParamDict): The parameters to be copied.

Returns:

ParamDict: Deep copies of the parameters.

nytorch.utils.copy_module(module: Module) → Module

Create a shallow copy of the module.

Args:

module (Module): The module to be copied.

Returns:

Module: The shallow copy of the module.

nytorch.utils.copy_modules(modules: OrderedDict[int, Module]) → OrderedDict[int, Module]

Create shallow copies of modules.

Args:

modules (ModuleDict): The modules to be copied.

Returns:

ModuleDict: Shallow copies of the modules.

nytorch.utils.make_modules_ref(modules: OrderedDict[int, Module], params: OrderedDict[int, Parameter], metas: OrderedDict[int, ModuleMeta], targets: Optional[Set[int]] = None) → None

Establish references between a set of modules and parameters according to correct structure.

Args:

modules (ModuleDict):

Modules needing to establish references.

params (ParamDict):

Parameters needing to establish references.

metas (MetaDict):

Records of the correct reference structure for modules and parameters.

targets (Optional[set[ModuleID]]):

Choose which modules to establish references for.

Examples:

class MyNet(nn.Module):
    def __init__(self):
        super().__init__()

def new_param(*args):
    return nn.Parameter(torch.randn(*args))

net0, net1 = MyNet(), MyNet()
param0, param1, param2 = new_param(1), new_param(2), new_param(3)

modules = OrderedDict([(0, net0), 
                       (1, net1)])
params = OrderedDict([(0, param0), 
                      (1, param1), 
                      (2, param2)])
metas = OrderedDict([(0, ModuleMeta(sub_modules=OrderedDict([("m1", 1)]), 
                                    sub_params=OrderedDict([("p0", 0)]))), 
                     (1, ModuleMeta(sub_modules=OrderedDict(), 
                                    sub_params=OrderedDict([("p1", 1), 
                                                            ("p2", 2)])))])

>>> make_modules_ref(modules, params, metas)
>>> assert net0.m1 is net1
>>> assert net0.p0 is param0
>>> assert net1.p1 is param1
>>> assert net1.p2 is param2

nytorch.version_updater module

class nytorch.version_updater.AddModuleVersionUpdater(module_id: int, attr_name: str, add_module: Module)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that adds a module to a particle.

This class facilitates the addition of a module to a particle by updating both version and particle data.

Args:

module_id (ModuleID):

ID of the module where the new module is added.

attr_name (str):

Name of the attribute where the new module is added.

add_module (Module):

The module to be added.

Attributes:

module_id (ModuleID):

ID of the module in the particle where the module is added.

attr_name (str):

Name of the attribute where the module is added.

add_module (Module):

The module to be added to the particle.

owner_modules (Optional[ModuleDict]):

Modules of the particle initiating the event.

owner_params (Optional[ParamDict]):

Parameters of the particle initiating the event.

Example:

class MyNet(NytoModule):
    ...

net0 = MyNet()
net1 = MyNet()
particle_kernel: ParticleKernel = net0._particle_kernel
particle_kernel.version_update(AddModuleVersionUpdater(net0._module_id, 
                                                       "attar_net1", 
                                                       net1))

>>> net0.attar_net1 is net1
True

add_module: Module

attr_name: str

module_id: int

owner_modules: Optional[OrderedDict[int, Module]]

owner_params: Optional[OrderedDict[int, Parameter]]

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

class nytorch.version_updater.AddParamVersionUpdater(module_id: int, attr_name: str, add_param: Parameter)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that adds a parameter to a particle.

Args:

module_id (ModuleID):

ID of the module where the new parameter is added.

attr_name (str):

Name of the attribute where the new parameter is added.

add_param (ParamType):

The parameter to be added.

Attributes:

module_id (ModuleID):

ID of the module in the particle where the parameter is added.

attr_name (str):

Name of the attribute where the parameter is added.

add_param (ParamType):

The parameter to be added to the particle.

owner_params (Optional[ParamDict]):

Parameters of the particle initiating the event.

Example:

class MyNet(NytoModule):
    ...

net = MyNet()
add_param = torch.nn.Parameter(torch.randn(3, 3))

particle_kernel: ParticleKernel = net._particle_kernel
particle_kernel.version_update(AddParamVersionUpdater(net._module_id, 
                                                      "attar_add_param", 
                                                      add_param))

>>> net.attar_add_param is add_param
True

add_param: Parameter

attr_name: str

module_id: int

owner_params: Optional[OrderedDict[int, Parameter]]

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

class nytorch.version_updater.DelBufferVersionUpdater(module_id: int, attr_name: str)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that deletes a buffer attribute from a particle.

Args:

module_id (ModuleID):

ID of the module performing this operation within the particle.

attr_name (str):

Name of the attribute being deleted.

Attributes:

module_id (ModuleID):

ID of the module in particle where this operation is executed.

attr_name (str):

Name of the attribute to be deleted.

Example:

class MyNet(NytoModule):
    ...

net = MyNet()
net.register_buffer("attar_buffer", torch.randn(3, 3))
particle_kernel: ParticleKernel = net._particle_kernel
particle_kernel.version_update(DelParamVersionUpdater(net._module_id, 
                                                      "attar_buffer"))

>>> hasattr(net, "attar_buffer")
False

attr_name: str

module_id: int

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

class nytorch.version_updater.DelModuleVersionUpdater(module_id: int, attr_name: str)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that deletes a module attribute from a particle.

Args:

module_id (ModuleID):

ID of the module performing this operation within the particle.

attr_name (str):

Name of the attribute being deleted.

Attributes:

module_id (ModuleID):

ID of the module in particle where this operation is executed.

attr_name (str):

Name of the attribute to be deleted.

Example:

class MyNet(NytoModule):
    ...

net = MyNet()
net.attar_moduel = MyNet()
particle_kernel: ParticleKernel = net._particle_kernel
particle_kernel.version_update(DelModuleVersionUpdater(net._module_id, 
                                                       "attar_moduel"))

>>> hasattr(net, "attar_moduel")
False

attr_name: str

module_id: int

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

class nytorch.version_updater.DelParamVersionUpdater(module_id: int, attr_name: str)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that deletes a parameter attribute from a particle.

Args:

module_id (ModuleID):

ID of the module performing this operation within the particle.

attr_name (str):

Name of the attribute being deleted.

Attributes:

module_id (ModuleID):

ID of the module in particle where this operation is executed.

attr_name (str):

Name of the attribute to be deleted.

Example:

class MyNet(NytoModule):
    ...

net = MyNet()
net.attar_param = torch.nn.Parameter(torch.randn(3, 3))
particle_kernel: ParticleKernel = net._particle_kernel
particle_kernel.version_update(DelParamVersionUpdater(net._module_id, 
                                                      "attar_moduel"))

>>> hasattr(net, "attar_param")
False

attr_name: str

module_id: int

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

class nytorch.version_updater.RegisterBufferVersionUpdater(module_id: int, attr_name: str, value: Optional[Tensor], persistent: bool)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that registers a buffer to a particle.

Args:

module_id (ModuleID):

ID of the module where the buffer is added.

attr_name (str):

Name of the attribute where the buffer is added.

value (Optional[torch.Tensor]):

The buffer to be added to the particle.

persistent (bool):

If True, the buffer becomes part of the module and is saved or loaded with it.

Attributes:

module_id (ModuleID):

ID of the module in the particle where the buffer is added.

attr_name (str):

Name of the attribute where the buffer is added.

value (Optional[torch.Tensor]):

The buffer to be added to the particle.

persistent (bool):

If True, the buffer becomes part of the module and is saved or loaded with it.

Example:

class MyNet(NytoModule):
    ...

net = MyNet()
add_tensor = torch.randn(3, 3)

particle_kernel: ParticleKernel = net._particle_kernel
particle_kernel.version_update(RegisterBufferVersionUpdater(net._module_id, 
                                                            "attar_add_tensor", 
                                                            add_tensor))

>>> net.attar_add_tensor is add_tensor
True

attr_name: str

module_id: int

persistent: bool

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

value: Optional[Tensor]

class nytorch.version_updater.SetModuleNoneVersionUpdater(module_id: int, attr_name: str)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that sets a module property in a particle to None.

Args:

module_id (ModuleID):

ID of the module performing this operation within the particle.

attr_name (str):

Name of the attribute being operated on.

Attributes:

module_id (ModuleID):

ID of the module performing this operation within the particle.

attr_name (str):

Name of the attribute being operated on.

Example:

class MyNet(NytoModule):
    ...

net = MyNet()
net.attar_module = MyNet()
particle_kernel: ParticleKernel = net._particle_kernel
particle_kernel.version_update(SetModuleNoneVersionUpdater(net._module_id, 
                                                           "attar_module"))

>>> net.attar_module is None
True

attr_name: str

module_id: int

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.

class nytorch.version_updater.SetParamNoneVersionUpdater(module_id: int, attr_name: str)

Bases: VersionUpdater[VersionDataImp, ParticleDataImp]

Updater for VersionKernel instances that sets a param property in a particle to None.

Args:

module_id (ModuleID):

ID of the module performing this operation within the particle.

attr_name (str):

Name of the attribute being operated on.

Attributes:

module_id (ModuleID):

ID of the module performing this operation within the particle.

attr_name (str):

Name of the attribute being operated on.

Example:

class MyNet(NytoModule):
    ...

net = MyNet()
net.attar_param = torch.nn.Parameter(torch.randn(3, 3))
particle_kernel: ParticleKernel = net._particle_kernel
particle_kernel.version_update(SetParamNoneVersionUpdater(net._module_id, 
                                                          "attar_param"))

>>> net.attar_param is None
True

attr_name: str

module_id: int

run(vdata: VersionDataImp) → tuple[VersionDataImp, ParticleUpdater]

Executes the update process.

This method generates an instance of the next version’s VersionData subclass and the current version’s ParticleUpdater.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Tuple[Tvdata, ParticleUpdater]: A tuple containing the next version’s VersionData instance and the current version’s ParticleUpdater.

set_particle_data(pdata: ParticleDataImp) → Self

Sets the particle data required before executing the update.

Args:

pdata (Tpdata): An instance of a ParticleData subclass that stores the current particle parameters.

Returns:

Self: The updated instance of the class.

set_version_data(vdata: VersionDataImp) → Self

Sets the version data required before executing the update.

Args:

vdata (Tvdata): An instance of a VersionData subclass that stores the current version’s particle metadata.

Returns:

Self: The updated instance of the class.