afnio.cognitive

class afnio.cognitive.Add

Bases: Module

Performs element-wise addition of two input Variables.

This module utilizes the Add operation from afnio.autodiff.basic_ops. The inputs must be instances of the Variable class. The forward method applies the addition operation to the .data field of the inputs and returns the resulting Variable.

Note

This module does not have any trainable parameters.

Example

>>> import afnio as hf
>>> from afnio import cognitive as cog
>>> class Addition(cog.Module):
...     def __init__(self):
...         super().__init__()
...         self.add = cog.Add()
>>>     def forward(self, x, y):
...         return self.add(x, y)
>>> input1 = hf.Variable(data="abc", role="input1")
>>> input2 = hf.Variable(data="def", role="input2")
>>> addition = Addition()
>>> result = addition(input1, input2)
>>> print(result)
'abcdef'
>>> print(result.role)
'input1 and input2'

Raises:

• TypeError – If either input is not an instance of Variable.

• TypeError – If addition between the input types is not allowed.

• ValueError – If a scalar .data is added to a list .data.

• ValueError – If list .data fields have mismatched lengths.

See also

afnio.autodiff.basic_ops.Add for the underlying operation.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x, y)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

class afnio.cognitive.ChatCompletion

Bases: Module

Generates a chat-based completion using a language model.

This module leverages the ChatCompletion operation from afnio.autodiff.lm_ops to perform model inference. The forward method accepts a list of messages representing the conversation history, with optional dynamic inputs for filling placeholders within the messages. The forward_model_client is responsible for interfacing with the language model (e.g., GPT), while completion_args allows customization of generation parameters such as temperature, maximum tokens, and seed.

Example

>>> import afnio as hf
>>> from afnio import cognitive as cog
>>> from afnio.models.openai import OpenAI
>>> from afnio import set_backward_model_client
>>> fwd_model_client = OpenAI()
>>> fwd_model_args = {"model": "gpt-4o", "temperature": 0.7}
>>> set_backward_model_client("openai/gpt-4o")
>>> class Assistant(cog.Module):
...     def __init__(self):
...         super().__init__()
...         self.chat = cog.ChatCompletion()
...     def forward(self, fwd_model, messages, inputs, **completion_args):
...         return self.chat(fwd_model, messages, inputs, **completion_args)
>>> system = Variable(
...     "You are a helpful assistant.",
...     role="system instruction",
...     requires_grad=True
... )
>>> user = Variable("Translate 'Hello' to {language}.", role="user query")
>>> language = hf.Variable("Italian", role="language")
>>> messages = [
...     {"role": "system", "content": [system]},
...     {"role": "user", "content": [user]},
... ]
>>> model = Assistant()
>>> response = model(
...     fwd_model_client,
...     messages,
...     inputs={"language": language},
...     **fwd_model_args
... )
>>> print(response.data)
'Ciao'
>>> feedback = Variable("Use only capital letters.", role="feedback")
>>> response.backward(feedback)
>>> system.grad[0].data
'The system instruction should enforce the use of capital letters only.'

See also

afnio.autodiff.lm_ops.ChatCompletion for the underlying operation.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_args: Dict[str, Any]

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(forward_model_client, messages, inputs=None, **completion_args)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

forward_model_client: Optional[ChatCompletionModel]

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

messages: List[Dict[str, List[Variable]]]

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

class afnio.cognitive.DeterministicEvaluator

Bases: Module

Evaluates predictions deterministically using a user-defined evaluation function.

This module utilizes the DeterministicEvaluator operation from afnio.autodiff.evaluator to compute evaluation scores and explanations. The forward method takes in a prediction, a target, an evaluation function (eval_fn), and its purpose description (eval_fn_purpose). It also accepts a reduction function (reduction_fn) and its purpose description (reduction_fn_purpose) to aggregate scores if needed. The method outputs an evaluation score and an explanation, both as Variable instances. The success_fn checks if all evaluations are successful, allowing the backward pass to skip unnecessary gradient computations. The method outputs an evaluation score and an explanation, both as Variable instances.

Example

>>> import afnio as hf
>>> from afnio import cognitive as cog
>>> from afnio import set_backward_model_client
>>> set_backward_model_client("openai/gpt-4o")
>>> class ExactColor(cog.Module):
...     def __init__(self):
...         super().__init__()
...         def exact_match_fn(pred: str, tgt: str) -> int:
...             return 1 if pred == tgt_data else 0
...         self.exact_match_fn = exact_match_fn
...         self.fn_purpose = "exact match"
...         self.reduction_fn = sum
...         self.reduction_fn_purpose = "summation"
...         self.exact_match = cog.DeterministicEvaluator()
...     def forward(self, prediction, target):
...         return self.exact_match(
...             prediction,
...             target,
...             self.exact_match_fn,
...             self.fn_purpose,
...             self.reduction_fn,
...             self.reduction_fn_purpose,
...         )
>>> prediction = hf.Variable(
...     data=["the color is green", "blue"],
...     role="color prediction",
...     requires_grad=True
... )
>>> target = ["green", "blue"]
>>> model = ExactColor()
>>> score, explanation = model(prediction, target)
>>> print(score.data)
1
>>> print(explanation.data)
'The evaluation function, designed for 'exact match', compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch, generating individual scores for each pair. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
>>> explanation.backward()
>>> prediction.grad[0].data
'Reassess the criteria that led to the initial prediction of 'green'.'

Raises:

TypeError – If inputs are not of the correct types.

See also

afnio.autodiff.evaluator.DeterministicEvaluator for the underlying operation.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

eval_fn: Callable[[Variable, Union[str, Variable]], List[Any]]

eval_fn_purpose: Union[str, Variable]

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(prediction, target, eval_fn, eval_fn_purpose, success_fn, reduction_fn, reduction_fn_purpose)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

reduction_fn: Optional[Callable[[List[Any]], Any]]

reduction_fn_purpose: Union[str, Variable, None]

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

success_fn: Optional[Callable[[List[Any]], bool]]

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

class afnio.cognitive.ExactMatchEvaluator

Bases: Module

Evaluates predictions using an exact match criterion.

This module leverages the ExactMatchEvaluator operation from afnio.autodiff.evaluator and is a specialized version of the DeterministicEvaluator that uses an exact matching function to compare the prediction and target. It returns an evaluation score (1 for exact match, 0 otherwise) and an explanation describing the evaluation result.

Example

>>> import afnio as hf
>>> from afnio import cognitive as cog
>>> from afnio import set_backward_model_client
>>> set_backward_model_client("openai/gpt-4o")
>>> class ExactColor(cog.Module):
...     def __init__(self):
...         super().__init__()
...         self.exact_match = cog.ExactMatchEvaluator()
...     def forward(self, prediction, target):
...         return self.exact_match(prediction, target)
>>> prediction = hf.Variable(
...     data="green",
...     role="color prediction",
...     requires_grad=True
... )
>>> target = "red"
>>> model = ExactColor()
>>> score, explanation = model(prediction, target)
>>> print(score.data)
0
>>> print(explanation.data)
'The evaluation function, designed for 'exact match', compared the <DATA> fields of the predicted variable and the target variable, resulting in a score: 0.'
>>> explanation.backward()
>>> system.grad[0].data
'Reassess the criteria that led to the initial prediction of 'green'.'

Raises:

TypeError – If inputs are not of the correct types.

See also

afnio.autodiff.evaluator.ExactMatchEvaluator for the underlying operation.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(prediction, target, reduction_fn=<built-in function sum>, reduction_fn_purpose='summation')

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

reduction_fn: Optional[Callable[[List[Any]], Any]]

reduction_fn_purpose: Union[str, Variable, None]

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

class afnio.cognitive.LMJudgeEvaluator

Bases: Module

Evaluates predictions using a language model (LM) as the judge.

This module leverages the LMJudgeEvaluator operation from afnio.autodiff.evaluator to perform model-based evaluations. The forward method accepts a list of messages that construct the evaluation prompt, with optional inputs to dynamically fill placeholders within message templates. A prediction is compared against a target (optional) to generate a score and an explanation.

When processing a batch of predictions and targets, reduction_fn function aggregates individual scores (e.g., using sum to compute a total score). The reduction_fn_purpose parameter is a brief description of the aggregation’s purpose (e.g., “summation”). If aggregation is not desired, set reduction_fn and reduction_fn_purpose to None. The success_fn checks if all evaluations are successful, allowing the backward pass to skip unnecessary gradient computations.

This module supports both evaluation (eval_mode=True) and optimization (eval_mode=False) modes.

The forward_model_client specifies the LM responsible for evaluation, while completion_args allows customization of generation parameters like temperature, max tokens, and seed.

Example

>>> import afnio as hf
>>> from afnio import cognitive as cog
>>> from afnio.models.openai import OpenAI
>>> from afnio import set_backward_model_client
>>> fwd_model_client = OpenAI()
>>> fwd_model_args = {"model": "gpt-4o", "temperature": 0.5}
>>> set_backward_model_client("openai/gpt-4o")
>>> class Evaluator(cog.Module):
...     def __init__(self):
...         super().__init__()
...         self.judge = cog.LMJudgeEvaluator()
...     def forward(self, fwd_model, messages, prediction, target, inputs, **completion_args):
...         return self.judge(fwd_model, messages, prediction, target, inputs, **completion_args)
>>> task = Variable(
...     "Evaluate if the translation is {metric}.",
...     role="evaluation task",
...     requires_grad=True
... )
>>> format = Variable(
...     "Provide 'score' (true/false) and 'explanation' in JSON.",
...     role="output format"
... )
>>> metric = Variable(["accurate", "accurate"], role="metric")
>>> user = Variable(
...     "<PREDICTION>{prediction}</PREDICTION><TARGET>{target}</TARGET>",
..      role="user query"
... )
>>> prediction = Variable(
...     ["Hola Mundo", "Salve a tutti"],
...     role="translated text",
...     requires_grad=True
... )
>>> target = ["Ciao Mondo", "Salve a tutti"]
>>> messages = [
...     {"role": "system", "content": [task, format]},
...     {"role": "user", "content": [user]},
... ]
>>> model = Evaluator()
>>> score, explanation = model(
...     fwd_model_client,
...     messages,
...     prediction,
...     target,
...     inputs={"metric": metric},
...     reduction_fn=sum,
...     reduction_fn_purpose="summation",
...     **fwd_model_args
... )
>>> print(score.data)
1
>>> print(explanation.data)
'The evaluation function, designed using an LM as the judge, compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
>>> explanation.backward()
>>> system.grad[0].data
'The translated text should be in Italian.'

See also

afnio.autodiff.evaluator.LMJudgeEvaluator for the underlying operation.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_args: Dict[str, Any]

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

eval_mode: Union[bool, Variable]

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(forward_model_client, messages, prediction, target=None, inputs=None, success_fn=None, reduction_fn=<built-in function sum>, reduction_fn_purpose='summation', eval_mode=True, **completion_args)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

forward_model_client: Optional[ChatCompletionModel]

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

messages: List[Dict[str, List[Variable]]]

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

reduction_fn: Optional[Callable[[List[Any]], Any]]

reduction_fn_purpose: Union[str, Variable, None]

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

success_fn: Optional[Callable[[List[Any]], bool]]

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

class afnio.cognitive.Module(*args, **kwargs)

Bases: object

Base class for all LM pipeline modules.

Your pipeline should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import afnio as hf
import afnio.cognitive as cog
import torch.cognitive.functional as F
from afnio.models.openai import OpenAI
from afnio import set_backward_model_client

fwd_model_client = OpenAI()
fwd_model_args = {"model": "gpt-4o", "temperature": 0.7}
set_backward_model_client("openai/gpt-4o")

class MedQA(cog.Module):
    def __init__(self):
        super().__init__()
        self.system_prompt = cog.Parameter(
            data="You are a doctor. Only answer medical questions on these areas:",
            role="system prompt",
            requires_grad=True,
        )
        self.topics = cog.Parameter(
            data="Dermatology and Cardiology",
            role="medical topics",
            requires_grad=False,
        )
        self.epilogue = hf.Variable(
            data="\nThank you for your query.",
            role="response preamble",
        )
        self.chat = cog.ChatCompletion()

    def forward(self, fwd_model, user_query, inputs, **completion_args):
        messages = [
            {"role": "system", "content": [self.system_prompt, self.topics]},
            {"role": "user", "content": [user_query]},
        ]
        response = self.chat(fwd_model, messages, inputs, **completion_args)
        return F.Add(response, self.epilogue)

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

• training (bool) – Boolean represents whether this module is in training or evaluation mode.

• automatic_optimization (bool) – Boolean that decides whether to use automatic optimization.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(*args, **kwargs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

class afnio.cognitive.Parameter(data=None, role=None, requires_grad=True)

Bases: Variable

A subclass of Variable that represents learnable parameters (similar to nn.Parameter). These parameters are typically text-based, learnable weights, embeddings, etc.

append_grad(gradient)

Appends a gradient value to the list .grad for this variable.

backward(gradient=None, retain_graph=None, create_graph=False, inputs=None)

Computes the gradient of current variable wrt graph leaves.

The graph is differentiated using the chain rule. If the variable is non-scalar (i.e. its data has more than one element) and requires gradient, the function additionally requires specifying a gradient. It should be a variable with data of matching type and shape, that represents the gradient of the differentiated function w.r.t. self.

This function accumulates gradients in the leaves - you might need to zero .grad attributes or set them to None before calling it.

Note

When inputs are provided, each input must be a leaf variable. If any input is not a leaf, a RuntimeError is raised.

Parameters:

• gradient (Variable, optional) – The gradient of the function being differentiated w.r.t. self. This argument can be omitted if self is a scalar.

• retain_graph (bool, optional) – If False, the graph used to compute the grads will be freed. Setting this to True retains the graph, allowing for additional backward calls on the same graph, useful for example for multi-task learning where you have multiple losses. However, retaining the graph is not needed in nearly all cases and can be worked around in a much more efficient way. Defaults to the value of create_graph.

• create_graph (bool, optional) – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products. Defaults to False.

• inputs (sequence of Variable, optional) – Inputs w.r.t. which the gradient will be accumulated into .grad. All other variables will be ignored. If not provided, the gradient is accumulated into all the leaf Variables that were used to compute the variables.

copy_(src)

Copies the data from the source Variable into this Variable.

Parameters:

src (Variable) – The source Variable to copy from.

Returns:

The current Variable with updated data, role and requires_grad.

Return type:

self

Raises:

• TypeError – If the source is not a Variable.

• ValueError – If the source data type does not match the target data type.

property data

detach()

Returns a new Variable, detached from the computation graph. This new Variable will not have a grad_fn and will not track gradients.

property grad: Variable | None

property grad_fn: Node | None

is_floating_point()

Checks if the Variable’s data contains floating-point values.

Returns:

True if the data is a floating-point type (either scalar or

all elements in a list/tuple are floating-point).

Return type:

bool

is_leaf: bool

All Variables that have requires_grad which is False will be leaf Variables by convention.

For Variables that have requires_grad which is True, they will be leaf Variables if they were created by the user. This means that they are not the result of an operation and so grad_fn is None.

Only leaf Variables will have their grad populated during a call to backward(). To get grad populated for non-leaf Variables, you can use retain_grad().

Example:

>>> a = hf.Variable("abc", requires_grad=True)
>>> a.is_leaf
True
>>> b = hf.Variable("abc", requires_grad=True).upper()
>>> b.is_leaf
False
# b was created by the operation that converts all string characters to uppercase
>>> c = hf.Variable("abc", requires_grad=True) + "def"
>>> c.is_leaf
False
# c was created by the addition operation
>>> d = hf.Variable("abc").upper()
>>> d.is_leaf
True
# d does not require gradients and so has no operation creating it (that is tracked by the autodiff engine)
>>> e = hf.Variable("abc").upper().requires_grad_()
>>> e.is_leaf
True
# e requires gradients and has no operations creating it

property output_nr: int

requires_grad: bool

requires_grad_(mode=True)

Change if autodiff should record operations on this variable: sets this variable’s requires_grad attribute in-place. Returns this variable.

requires_grad_()’s main use case is to tell autodiff to begin recording operations on a Variable variable. If variable has requires_grad=False (because it was obtained through a DataLoader, or required preprocessing or initialization), variable.requires_grad_() makes it so that autodiff will begin to record operations on variable.

Parameters:

requires_grad (bool) – If autodiff should record operations on this variable. Default: True.

Example

>>> # Initialize with requires_grad=False for data preprocessing
>>> x = hf.Variable(data="abc", role="input")
>>> x = preprocess(x)  # Preprocess without gradient tracking
>>> x
variable(abc, role=input, requires_grad=False)

>>> # Now enable requires_grad for backpropagation
>>> x.requires_grad_()
>>> output = model(x)
>>> output.backward()  # Backpropagation through `x`
>>> x.grad
variable(ABC, role=input, requires_grad=True)

retain_grad()

Enable gradient retention for non-leaf variables.

to(dtype=None)

Cast the data of the Variable to the specified dtype.

Parameters:

dtype (Optional[type]) – The target type to cast the data (e.g., float, int, str).

Returns:

A new Variable with data cast to the target dtype.

Return type:

Variable

variable_id: Optional[str]

class afnio.cognitive.Split

Bases: Module

Splits a single input Variable into multiple output Variables.

This module utilizes the Split operation from afnio.autodiff.basic_ops. It supports string data types, splitting the string data of the input Variable based on a specified delimiter and an optional maximum number of splits.

Note

This module does not have any trainable parameters.

Example

>>> import afnio as hf
>>> from afnio import cognitive as cog
>>> class Splitter(cog.Module):
...     def __init__(self):
...         super().__init__()
...         self.split = cog.Split()
>>>     def forward(self, x):
...         return self.split(x, " ", 1)
>>> input = hf.Variable(data="Afnio is great!", role="sentence")
>>> splitter = Splitter()
>>> result = splitter(input)
>>> print([r.data for r in result])
['Afnio', 'is great!']
>>> print([r.role for r in result])
['split part 0 of sentence', 'split part 1 of sentence']

Raises:

• TypeError – If the input is not an instance of Variable.

• TypeError – If the input Variable’s data is not a string.

See also

afnio.autodiff.basic_ops.Split for the underlying operation.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x, sep=None, maxsplit=-1)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

maxsplit: Union[int, Variable, None]

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

sep: Union[str, Variable, None]

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

class afnio.cognitive.Sum

Bases: Module

Aggregates a list of input Variables into a single output Variable.

This module utilizes the Sum operation from afnio.autodiff.basic_ops. It supports both numerical (int, float) and string data types. For numerical data, it computes the sum. For string data, it concatenates the values and wraps each in <ITEM></ITEM> tags.

Note

This module does not have any trainable parameters.

Example

>>> import afnio as hf
>>> from afnio import cognitive as cog
>>> class Summation(cog.Module):
...     def __init__(self):
...         super().__init__()
...         self.sum = cog.Sum()
>>>     def forward(self, x):
...         return self.sum(x)
>>> input1 = hf.Variable(data="abc", role="input1")
>>> input2 = hf.Variable(data="def", role="input2")
>>> input3 = hf.Variable(data="ghi", role="input3")
>>> summation = Summation()
>>> result = summation([input1, input2, input3])
>>> print(result.data)
'<ITEM>abc</ITEM><ITEM>def</ITEM><ITEM>ghi</ITEM>'
>>> print(result.role)
'input1 and input2 and input3'

Raises:

TypeError – If any input is not an instance of Variable.

See also

afnio.autodiff.basic_ops.Sum for the underlying operation.

T_destination = ~T_destination

automatic_optimization: bool

buffers(recurse=True)

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

hf.Variable – module buffer

Example:

>>> for buf in model.buffers():
>>>     print(type(buf), buf.data)
<class 'afnio.Variable'> ("Structure your answer as JSON.")
<class 'afnio.Variable'> ("Use the format\n\n{\n  \"response\": \"Your concise answer here.\"\n}")

chats(recurse=True)

Return an iterator over module multi-turn chats.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

Yields:

MultiTurnMessages – module chats

Example:

>>> for chat in pipeline.chats():
>>>     print(type(chat), chat)
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a doctor., role=system instruction, requires_grad=False)]}, {'role': 'user', 'content': [Variable(data=Is {item} a disease?, role=user query, requires_grad=False)]}]
<class 'cog.MultiTurnMessages'> [{'role': 'system', 'content': [Variable(data=You are a helpful assistant., role=system instruction, requires_grad=False), Variable(data=Only answer with YES or NO., role=user query, requires_grad=False)]}]

children()

Return an iterator over immediate children modules.

Yields:

Module – a child module

completion_configs(recurse=True)

Return an iterator over registered completion configs.

Parameters:

recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

Yields:

dict – completion arguments

Example::

>>> for config in model.completion_configs():
>>>     print(config)
{"model": "gpt-4o", "seed": 42, "temperature": 0}

configure_optimizers()

Configure and return the optimizer for this module.

This method should be implemented in subclasses to define the optimizer configuration. It is called by the Trainer to set up the optimization routine.

Returns:

An instance of an optimizer configured for this module.

Return type:

Optimizer

Raises:

NotImplementedError – If not implemented in a subclass.

empty_grad()

Reset gradients of all model parameters and content variables in chats’ messages.

This method is useful for clearing out gradients before starting a new optimization step. It ensures that both module parameters and Variables within multi-turn chat’s message contents have their gradients reset, avoiding unintended gradient accumulation.

eval()

Set the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

This is equivalent with self.train(False).

Returns:

self

Return type:

Module

abstractmethod extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

One should invoce the Module instance (Module.__call__ method) instead of directly calling Module.forward(). This way hooks are registered and run.

functions(recurse=True)

Return an iterator over registered functions.

Parameters:

recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

Yields:

Callable – functions

Example::

>>> for func in model.functions():
>>>     print(func)
<built-in function sum>
<function my_func at 0x7e7a0665b9c0>

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict.

Returns:

Any extra state to store in the module’s state_dict.

Return type:

object

load_state_dict(state_dict, strict=True, assign=False, model_clients=None)

Copy parameters, buffers, chats, models, completion configs and functions from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict.

Parameters:

• state_dict (dict) – A dict containing parameters, persistent buffers, chats, models, completion configs and functions.

• strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When False, the properties of the Variables in the current module are preserved while when True, the properties of the Variables in the state dict are preserved. The only exception is the requires_grad field of Default: ``False`

• model_clients (dict, optional) – A dictionary mapping model client keys (e.g., ‘fw_model_client’) to their respective instances of BaseModel. These instances will be used to reconstruct any model clients referenced within the optimizer state. If a required model client is missing, an error will be raised with instructions on how to provide the missing client.

Returns:

• missing_keys is a list of str containing any keys that are

  expected by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter, or buffer, or chat, or model, or completion config, or function is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

models(recurse=True)

Return an iterator over module language model clients.

Parameters:

recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

Yields:

BaseModel – module model

Example:

>>> for model in pipeline.models():
>>>     print(type(model))
<class 'afnio.models.openai.AsyncOpenAI'>

modules()

Return an iterator over all modules in the network.

Yields:

Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...         super().__init__()
...         add = cog.Add()
...         self.module1 = add
...         self.module2 = add
>>>     def forward(self, x, y):
...         out1 = self.module1(x, x)
...         out2 = self.module2(x, y)
...         return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.modules()):
...     print(idx, '->', m)
0 -> MyModel(
(module1): Module()
(module2): Module()
)
1 -> Module()

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

• prefix (str) – prefix to prepend to all buffer names.

• recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

• remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, hf.Variable) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>     if "format_type" in name:
>>>         print(param.data)

named_chats(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module multi-turn chats, yielding both the name of chat as well as the chat itself.

Parameters:

• prefix (str) – prefix to prepend to all chat names.

• recurse (bool) – if True, then yields chats of this module and all submodules. Otherwise, yields only chats that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated chats in the result. Defaults to True.

Yields:

(str, MultiTurnMessages) – Tuple containing the name and chat

Example:

>>> for name, chat in self.named_chats():
>>>     if "messages" in name:
>>>         print(messages[0]["role"])

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

named_completion_configs(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module completion configs, yielding both the name of the completion config as well as the completion config itself.

Parameters:

• prefix (str) – prefix to prepend to all completion config names.

• recurse (bool) – if True, then yields completion configs of this module and all submodules. Otherwise, yields only completion configs that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated completion configs in the result. Defaults to True.

Yields:

(str, dict) – Tuple containing the name and completion configs

Example:

>>> for name, config in self.named_completion_configs():
>>>     print(name, type(config))
chat.completion_args {'model': 'gpt-4o', 'seed': 42, 'temperature': 0}

named_functions(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module functions, yielding both the name of the function as well as the function itself.

Parameters:

• prefix (str) – prefix to prepend to all function names.

• recurse (bool) – if True, then yields functions of this module and all submodules. Otherwise, yields only functions that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated functions in the result. Defaults to True.

Yields:

(str, Callable) – Tuple containing the name and functions

Example:

>>> for name, func in self.named_functions():
>>>     print(name, func)
reduction_fn <built-in function sum>
eval_fn <function my_func at 0x7e7a0665b9c0>

named_models(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module model clients, yielding both the name of the model as well as the model itself.

Parameters:

• prefix (str) – prefix to prepend to all model names.

• recurse (bool) – if True, then yields models of this module and all submodules. Otherwise, yields only models that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated models in the result. Defaults to True.

Yields:

(str, BaseModel) – Tuple containing the name and model

Example:

>>> for name, model in self.named_models():
>>>     print(name, type(model))
model_client <class 'afnio.models.openai.AsyncOpenAI'>

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

• memo (Optional[Set[Module]]) – a memo to store the set of modules already added to the result

• prefix (str) – a prefix that will be added to the name of the module

• remove_duplicate (bool) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, add will be returned only once.

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules()):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())

Example:

>>> class MyPipeline(cog.Module):
...     def __init__(self):
...     super().__init__()
...     add = cog.Add()
...     self.module1 = add
...     self.module2 = add
>>> def forward(self, x, y):
...     out1 = self.module1(x, x)
...     out2 = self.module2(x, y)
...     return out1 + out2
>>> pipeline = MyPipeline()
>>> for idx, m in enumerate(model.named_modules(remove_duplicate=False)):
...     print(idx, '->', m)
0 -> ('', MyModel(
(module1): Module()
(module2): Module()
))
1 -> ('module1', Module())
2 -> ('module2', Module())

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>     if "prompt" in name:
>>>         print(param.data)

optimizers()

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

This method is useful for accessing the optimizer(s) configured in the configure_optimizers() method by the fit() method.

Returns:

The optimizer(s) used by this module.

Return type:

Union[Optimizer, List[Optimizer]]

Example:

>>> optimizers = model.optimizers()
>>> for optimizer in optimizers:
>>>     print(optimizer)
TGD (
Parameter Group 0
    completion_args: {'model': 'gpt-4.1'}
    constraints: []
    inputs: {}
    messages: [
    {'role': 'system', 'content': [Variable(data="Placeholder Textual Gradient Descent optimizer system prompt", role=Textual Gradient Descent optimizer system prompt, requires_grad=False)]},
    {'role': 'user', 'content': [Variable(data="Placeholder for Textual Gradient Descent optimizer user prompt", role=Textual Gradient Descent optimizer user prompt, requires_grad=False)]}
    ]
    model_client: <afnio.models.openai.AsyncOpenAI object at 0x710df9c149a0>
    momentum: 3
)

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> for param in pipeline.parameters():
>>>     print(type(param), param.data)
<class 'cog.Parameter'> ("You are a doctor.")
<class 'cog.Parameter'> ("Only answer with YES or NO.")

register_buffer(name, variable, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, Prompt’s format_type 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.

Parameters:

• name (str) – Name of the buffer. The buffer can be accessed from this module using the given name.

• variable (Variable or None) – Buffer to be registered. If None, then operations that run on buffers are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – Whether the buffer is part of this module’s state_dict.

Example::

>>> self.register_buffer('format_type', hf.Variable(data="Structure your answer as JSON.", role="format type"))

register_chat(name, messages)

Add multi-turn chat messages to the module.

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

Parameters:

• name (str) – Name of the chat. The chat can be accessed from this module using the given name.

• messages (MultiTurnMessages or None) – Chat to be added to the module. If None, then operations that run on chats are ignored. If None, the chat is not included in the module’s state_dict.

register_completion_config(name, args)

Register completion-specific arguments for text generation.

This method allows dynamic storage of completion-related parameters such as temperature, max_tokens, top_p, etc.

Parameters:

• name (str) – Name of the completion argument set.

• args (dict or None) – Dictionary of completion arguments. If None, the argument is not included in the module’s state_dict.

register_function(name, func)

Add a function to the module.

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

Parameters:

• name (str) – Name of the function. The function can be accessed from this module using the given name.

• func (FunctionType or None) – A standard Python function (i.e., a def-defined function, not a lambda or callable object) that can be pickled and registered for later execution. If None, the function is unregistered. If None, the function is not included in the module’s state_dict.

register_model(name, model)

Add language model the module.

The language model can be accessed as an attribute using given name.

Parameters:

• name (str) – Name of the model. The model can be accessed from this module using the given name.

• model (BaseModel or None) – Model to be added to the module. If None, then operations that run on models are ignored. If None, the model is not included in the module’s state_dict.

register_module(name, module)

Add a child module to the current module.

This method explicitly adds a child module to the current module’s hierarchy. The child module can then be accessed as an attribute using the given name and will be registered in the _modules dictionary.

When to use: - Use register_module() when dynamically adding submodules at runtime, especially when the submodule name is determined programmatically. - This can be useful for creating flexible and modular architectures.

When it’s unnecessary: - Directly assigning the module to an attribute (e.g., self.module_name = SubModule()) automatically registers it, so using register_module() is unnecessary in such cases.

Parameters:

• name (str) – 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.

Raises:

• TypeError – If module is not a subclass of Module or if name is not a string.

• KeyError – If name is already an attribute of the module but not in _modules, or if name contains invalid characters such as ‘.’ or is empty.

Example::

>>> class DynamicPipeline(cog.Module):
>>>     def __init__(self):
>>>         super().__init__()
>>>         # Dynamically add submodules
>>>         for i in range(3):
>>>             self.register_module(f"layer_{i}", cog.Module())

>>> pipeline = DynamicPipeline()
>>> print(pipeline._modules.keys())
odict_keys(['layer_0', 'layer_1', 'layer_2'])

Note

If assigning submodules using standard attribute assignment (e.g., self.submodule = SubModule()), calling register_module() explicitly is not required. Direct assignment automatically registers the module.

register_parameter(name, param)

Add a parameter to the module.

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

Parameters:

• name (str) – Name of the parameter. The parameter can be accessed from this module using the given name.

• param (Parameter or None) – Parameter to be added to the module. If None, then operations that run on parameters are ignored. If None, the parameter is not included in the module’s state_dict.

requires_grad_(requires_grad=True)

Change if autodiff should record operations on parameters and chats in this module.

This method sets the requires_grad attributes of all module parameters in-place. It also sets the requires_grad attributes of all the Variables within the content of multi-turn chats.

Effect on Parameters:

• Sets requires_grad for each registered parameter in the module.

Effect on Chats:

• Iterates through all multi-turn chats and sets requires_grad

for each Variable in the “content” key of the chat’s message.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually.

Parameters:

requires_grad (bool) – Whether autodiff should record operations on parameters and chats in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict.

state_dict(*, destination=None, prefix='', keep_vars=False)

Return a dictionary containing references to the whole state of the module.

Parameters, persistent buffers (e.g. running averages), multi-turn chats, models, completion configs and functions are included. Keys are corresponding parameter, buffer, chat, model, completion config and function names. Parameters, buffers, chats, models, completion configs and functions set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters, buffers, chats, models, completion configs and functions.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – A prefix added to parameter, buffer, chat, model, completion config and function names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – By default the Variable s returned in the state dict are detached from autodiff. If it’s set to True, detaching will not be performed. Default: False.

Returns:

A dictionary containing a whole state of the module.

Return type:

dict

Example:

>>> module.state_dict().keys()
['system_prompt', 'classification_labels', 'format_type', 'user_prompt']

test_step(batch, batch_idx)

Perform a single test step.

This method should be implemented in subclasses to define the test logic. It is called by the Trainer during the testing loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

training: bool

training_step(batch, batch_idx)

Perform a single training step.

This method should be implemented in subclasses to define the training logic. It is called by the Trainer during the training loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

validation_step(batch, batch_idx)

Perform a single validation step.

This method should be implemented in subclasses to define the validation logic. It is called by the Trainer during the validation loop.

Parameters:

• batch (Any) – The output of your data iterable, normally a DataLoader.

• batch_idx (int) – The index of this batch.

Returns:

The loss as a tuple of two Variables:

• The evaluation score (a Variable containing the loss value).

• The explanation (a Variable containing a string explanation of the evaluation result).

• dict: A dictionary. Can include any keys, but must include

  the key 'loss' containing a tuple of two Variables (score and explanation).

• None: Skip to the next batch.

Return type:

• Tuple[Variable, Variable]

Raises:

NotImplementedError – If not implemented in a subclass.

afnio.cognitive.add(x, y)

Implements an addition operation for Variable instances within the afnio framework, supporting automatic differentiation.

The Add function supports both scalar and list .data fields:

• Scalars: Adds numerical values (int, float) or concatenates strings.

• Lists: Performs element-wise addition of corresponding elements from the lists. Lists must be of the same length.

It automatically handles type-based operations:

• For numerical data (int, float), it performs arithmetic addition.

• For strings, it concatenates the values.

• Mixed types (e.g., string and number) are converted appropriately before performing the addition.

This operation also tracks Variable dependencies, enabling automatic gradient computation through backpropagation.

Example with scalar inputs:

>>> x = Variable(data="abc", role="first input", requires_grad=True)
>>> y = Variable(data="def", role="second input", requires_grad=False)
>>> result = F.add(x, y)
>>> result.data
'abcdef'
>>> result.role
'first input and second input'
>>> result.requires_grad
True
>>> g = Variable(data="MY_FEEDBACK", role="add gradient")
>>> result.backward(g)
>>> x.grad.data
'Here is the combined feedback we got for this specific first input and other variables: MY_FEEDBACK'
>>> x.grad.role
'feedback to first input'

Example with batched inputs:

>>> x = Variable(data=[1, 2, 3], role="first input", requires_grad=True)
>>> y = Variable(data=[4, 5, 6], role="second input", requires_grad=False)
>>> result = F.add(x, y)
>>> result.data
[5, 7, 9]
>>> result.role
'first input and second input'
>>> result.requires_grad
True

afnio.cognitive.chat_completion(forward_model_client, messages, inputs=None, **completion_args)

Implements a chat completion operation using the specified language model within the afnio framework, supporting automatic differentiation.

Features:

• Mini-Batching: Processes multiple input dictionaries simultaneously to improve

  throughput.

• Asynchronous Execution: Both the forward and backward passes are optimized to

  run asynchronous calls for each mini-batch, reducing latency.

• Gradient Computation: Supports automatic differentiation for all Variables

  in messages and inputs arguments, maintaining the order of gradients.

The ChatCompletion function generates a Variable responses by passing a composite prompt, built from a list of messages and optional inputs, to the forward_model_client. Each message is a dictionary with a ‘role’ (e.g., ‘system’, ‘user’) and a list of Variable objects as ‘content’. inputs is a dictionary containing strings, list of strings or Variable``s providing dynamic values to fill placeholders within message templates. If ``inputs contain lists of strings or Variable``s which ``.data field is a list, the response’s .data field will be a list, corresponding to the batched results. Otherwise, the .data field will be a scalar string. Additional behavior, such as temperature or token limits, can be customized through completion_args.

Example with scalar inputs:

>>> system = Variable(
...     "You are a helpful assistant.",
...     role="system instruction",
...     requires_grad=True
... )
>>> user = Variable("Translate 'Hello' to {language}.", role="user query")
>>> messages = [
...     {"role": "system", "content": [system]},
...     {"role": "user", "content": [user]},
... ]
>>> inputs = {"language": Variable("Italian", role="language")}
>>> response = F.chat_completion(
...     model_client,
...     messages,
...     inputs=inputs,
...     temperature=0.7
... )
>>> print(response.data)
'Ciao'
'Hola'
>>> feedback = Variable("Use only capital letters.", role="feedback")
>>> response.backward(feedback)
>>> system.grad[0].data
'The system instruction should enforce the use of capital letters only.'

Example with batched inputs:

>>> system = Variable(
...     "You are a helpful assistant.",
...     role="system instruction",
...     requires_grad=True
... )
>>> user = Variable("Translate 'Hello' to {language}.", role="user query")
>>> messages = [
...     {"role": "system", "content": [system]},
...     {"role": "user", "content": [user]},
... ]
>>> inputs = {
...     "language": [
...         Variable("Italian", role="language"),
...         Variable("Spanish", role="language")
...     ]
... }
>>> response = F.chat_completion(
...     model_client,
...     messages,
...     inputs=inputs,
...     temperature=0.7
... )
>>> print(response.data)
['Ciao', 'Hola']

afnio.cognitive.deterministic_evaluator(prediction, target, eval_fn, eval_fn_purpose, success_fn, reduction_fn, reduction_fn_purpose)

Evaluates predictions deterministically using a user-defined evaluation function within the afnio framework, supporting automatic differentiation.

The DeterministicEvaluator function computes a score and an explanation based on the prediction and target inputs using a user-defined evaluation function (eval_fn). The evaluation function’s purpose is described by eval_fn_purpose. Outputs include a numerical or textual score and a textual explanation, both wrapped as Variable objects.

The prediction is a Variable. The target can be a string, a list of strings, or a Variable. Each Variable passed as an input argument can have either a scalar or a list .data field, supporting both individual samples and batch processing. For batch processing, the lengths of prediction and target must match.

The success_fn parameter is a user-defined function that returns True when all predictions evaluated by eval_fn are considered successful, and False otherwise. If success_fn returns True, the backward pass will skip gradient calculations and directly return an empty gradient, optimizing computational time.

The reduction_fn parameter specifies the aggregation function to use for scores across a batch of predictions and targets. When specified, the reduction function’s purpose is described using reduction_fn_purpose. If aggregation is not desired, set reduction_fn and reduction_fn_purpose to None.

Example with scalar inputs:

>>> prediction = Variable(
...     data="green",
...     role="color prediction",
...     requires_grad=True
... )
>>> target = "red"
>>> def exact_match_fn(p: str, t: str) -> int:
...     return 1 if p == t else 0
>>> score, explanation = F.deterministic_evaluator(
...     prediction,
...     target,
...     exact_match_fn,
...     "exact match",
... )
>>> score.data
0
>>> explanation.data
'The evaluation function, designed for 'exact match', compared the <DATA> field of the predicted variable ('green') with the <DATA> field of the target variable ('red'), resulting in a score: 0.'
>>> explanation.backward()
>>> prediction.grad[0].data
'Reassess the criteria that led to the initial prediction of 'green'.'

Example with batched inputs:

>>> prediction = Variable(
...     data=["green", "blue"],
...     role="color prediction",
...     requires_grad=True
... )
>>> target = ["red", "blue"]
>>> def exact_match_fn(p: str, t: str) -> int:
...     return 1 if p == t else 0
>>> score, explanation = F.deterministic_evaluator(
...     prediction,
...     target,
...     exact_match_fn,
...     "exact match",
...     reduction_fn=sum,
...     reduction_fn_purpose="summation"
... )
>>> score.data
1
>>> explanation.data
'The evaluation function, designed for 'exact match', compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch, generating individual scores for each pair. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
>>> explanation.backward()
>>> prediction.grad[0].data
'Reassess the criteria that led to the initial prediction of 'green'.'

afnio.cognitive.exact_match_evaluator(prediction, target, reduction_fn=<built-in function sum>, reduction_fn_purpose='summation')

Evaluates predictions using exact matching within the afnio framework, supporting automatic differentiation.

The ExactMatchEvaluator function computes a score and an explanation by comparing the data fields of a prediction and a target for an exact match. For each sample:

• A score of 1 is assigned for an exact match.

• A score of 0 is assigned otherwise.

The prediction is a Variable. The target can be a string, a list of strings, or a Variable. Each Variable passed as an input argument can have either a scalar or a list .data field, supporting both individual samples and batch processing. For batch processing, the lengths of prediction and target must match.

If batched inputs are provided, the scores can be aggregated using an optional reduction_fn, such as sum. The purpose of the reduction is described using reduction_fn_purpose. If aggregation is not desired, set reduction_fn and reduction_fn_purpose to None.

Example with scalar inputs:

>>> prediction = Variable(
...     data="green",
...     role="color prediction",
...     requires_grad=True
... )
>>> target = "red",
>>> score, explanation = F.exact_match_evaluator(prediction, target)
>>> score.data
0
>>> explanation.data
'The evaluation function, designed for 'exact match', compared the <DATA> field of the predicted variable ('green') with the <DATA> field of the target variable ('red'), resulting in a score: 0.'
>>> explanation.backward()
>>> prediction.grad[0].data
'Reassess the criteria that led to the initial prediction of 'green'.'

Example with batched inputs:

>>> prediction = Variable(
...     data=["green", "blue"],
...     role="color prediction",
...     requires_grad=True
... )
>>> target = ["red", "blue"]
>>> score, explanation = F.exact_match_evaluator(prediction, target)
>>> score.data
1
>>> explanation.data
'The evaluation function, designed for 'exact match', compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch, generating individual scores for each pair. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
>>> explanation.backward()
>>> prediction.grad[0].data
'Reassess the criteria that led to the initial prediction of 'green'.'

afnio.cognitive.lm_judge_evaluator(forward_model_client, messages, prediction, target=None, inputs=None, success_fn=None, reduction_fn=<built-in function sum>, reduction_fn_purpose='summation', eval_mode=True, **completion_args)

Implements an evaluation of a model prediction using a language model (LM) as the judge within the afnio framework, supporting automatic differentiation.

This function returns a score and an explanation, both as Variable objects, by comparing a prediction against a target (when present) using a composite prompt. The prompt is constructed from a list of messages and optional inputs, which can dynamically populate placeholders in the message templates. The evaluation process leverages the specified forward_model_client to perform the LM-based assessment.

The prediction is a Variable. The target can be a string, a list of strings, or a Variable. Similarly, the inputs dictionary can include strings, lists of strings, or Variable``s. Each ``Variable passed as an input argument can have either a scalar or a list .data field, supporting both individual samples and batch processing. For batch processing, the lengths of prediction, target, and any batched inputs must match.

The success_fn parameter is a user-defined function that returns True when all predictions evaluated by the LM as Judge are considered successful, and False otherwise. If success_fn returns True, the backward pass will skip gradient calculations and directly return an empty gradient, optimizing computational time.

If you are processing a batch of predictions and targets, you can use the reduction_fn to aggregate individual scores (e.g., using sum to compute a total score). The reduction_fn_purpose parameter is a brief description of the aggregation’s purpose (e.g., “summation”). If you don’t want any aggregation, set both reduction_fn and reduction_fn_purpose to None.

The function operates in two modes controlled by eval_mode:

• ``eval_mode=True`` (default) – Computes gradients for prediction only. Use it for direct feedback on predictions.

• ``eval_mode=False`` – Computes gradients for messages and inputs. Use it to optimize the evaluator or align with human evaluation datasets.

Additional model parameters, such as temperature, max tokens, or seed values, can be passed through completion_args to customize the LLM’s behavior.

Example with scalar inputs:

>>> task = Variable(
...     "Evaluate if the translation is accurate.",
...     role="evaluation task",
...     requires_grad=True
... )
>>> format = Variable(
...     "Provide 'score' (true/false) and 'explanation' in JSON.",
...     role="output format"
... )
>>> user = Variable(
...     "<PREDICTION>{prediction}</PREDICTION><TARGET>{target}</TARGET>",
...     role="user query"
... )
>>> prediction = Variable(
...     "Hola Mundo",
...     role="translated text",
...     requires_grad=True
... )
>>> target = Variable("Ciao Mondo", role="expected output")
>>> messages = [
...     {"role": "system", "content": [task, format]},
...     {"role": "user", "content": [user]}
... ]
>>> score, explanation = F.lm_judge_evaluator(
...     model,
...     messages,
...     prediction,
...     target,
...     temperature=0.5,
... )
>>> score.data
False
>>> explanation.data
'The translated text is in Spanish, but the expected is in Italian.'
>>> explanation.backward()
>>> prediction.grad[0].data
'The translated text should be in Italian.'

Example with batched inputs:

>>> task = Variable(
...     "Evaluate if the translation is accurate.",
...     role="evaluation task",
...     requires_grad=True
... )
>>> format = Variable(
...     "Provide 'score' (true/false) and 'explanation' in JSON.",
...     role="output format"
... )
>>> user = Variable(
...     "<PREDICTION>{prediction}</PREDICTION><TARGET>{target}</TARGET>",
...     role="user query"
... )
>>> prediction = Variable(
...     data=["Hola Mundo", "Salve a tutti"],
...     role="translated text",
...     requires_grad=True,
... )
>>> target = ["Ciao Mondo", "Salve a tutti"]
>>> score, explanation = F.lm_judge_evaluator(
...     model,
...     messages,
...     prediction,
...     target,
...     reduction_fn=sum,
...     reduction_fn_purpose="summation",
... )
>>> score.data
1
>>> explanation.data
'The evaluation function, designed using an LM as the judge, compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
>>> explanation.backward()
>>> prediction.grad[0].data
'The translated text should be in Italian.'

afnio.cognitive.split(x, sep=None, maxsplit=-1)

Implements a split operation for Variable instances within the afnio framework, supporting automatic differentiation.

The Split function divides the .data of the input Variable into multiple parts using a specified delimiter sep. If maxsplit is specified, the split operation is limited to a maximum number of splits. It handles both scalar and list .data fields:

• Scalars: The scalar .data (a single string) is split into substrings based on the specified sep and maxsplit parameters.

• Lists: Each element of the list .data (strings) is split individually. If splits of varying lengths occur, shorter splits are automatically padded with empty strings to ensure consistent dimensions.

During backpropagation, feedback is collected and aggregated across all split parts. The combined feedback is propagated back to the original input Variable, allowing for the proper computation of gradients.

Example with scalar inputs:

>>> x = Variable(data="afnio is great!", role="sentence", requires_grad=True)
>>> result = Split.apply(x, sep=" ", maxsplit=1)
>>> [var.data for var in result]
['afnio', 'is great!']
>>> result[0].role
'split part 0 of sentence'
>>> g_1 = Variable(data="MY_FIRST_FEEDBACK", role="gradient")
>>> g_2 = Variable(data="MY_SECOND_FEEDBACK", role="gradient")
>>> result[0].backward(g_1, retain_graph=True)
>>> result[1].backward(g_2)
>>> x.grad[0].data
'Here is the combined feedback we got for this specific sentence and other variables: <ITEM>MY_FIRST_FEEDBACK</ITEM><ITEM></ITEM>'
>>> x.grad[0].role
'feedback to sentence'
>>> x.grad[1].data
'Here is the combined feedback we got for this specific sentence and other variables: <ITEM></ITEM><ITEM>MY_SECOND_FEEDBACK</ITEM>'
>>> x.grad[1].role
'feedback to sentence'

Example with batched inputs:

>>> x = Variable(
...     data=["afnio is great!", "Deep learning"],
...     role="sentences",
...     requires_grad=True
... )
>>> result = Split.apply(x, sep=" ", maxsplit=2)
>>> [var.data for var in result]
[['afnio', 'Deep'], ['is', 'learning'], ['great!', '']]
>>> g = Variable(data="MY_FEEDBACK", role="gradient")
>>> result[1].backward(g)
>>> x.grad[0].data
'Here is the combined feedback we got for this specific sentences and other variables: <ITEM></ITEM><ITEM>MY_FEEDBACK</ITEM><ITEM></ITEM>'
>>> x.grad[0].role
'feedback to sentences'

afnio.cognitive.sum(x)

Implements a summation operation for a list of Variable instances within the afnio framework, supporting automatic differentiation.

The Sum function aggregates the .data, .role, and .requires_grad attributes of all input Variable instances into a single Variable. It supports both scalar and list .data fields:

• Scalars: Computes the arithmetic sum for numerical data (int, float) or concatenates all string values, wrapping each in <ITEM></ITEM> tags.

• Lists: Aggregates the corresponding elements of the lists. For numerical data, it sums the corresponding elements. For string data, it concatenates them, wrapping each element in <ITEM></ITEM> tags.

During backpropagation, the function distributes the gradient to all input Variable instances that require gradients.

Example with scalar inputs:

>>> x = Variable(data="abc", role="first input", requires_grad=True)
>>> y = Variable(data="def", role="second input", requires_grad=False)
>>> result = F.sum([x, y])
>>> result.data
'<ITEM>abc</ITEM><ITEM>def</ITEM>'
>>> result.role
'first input and second input'
>>> result.requires_grad
True
>>> g = Variable(data="MY_FEEDBACK", role="add gradient")
>>> result.backward(g)
>>> x.grad.data
'Here is the combined feedback we got for this specific first input and other variables: MY_FEEDBACK'
>>> x.grad.role
'feedback to first input'

Example with batched inputs:

>>> x = Variable(data=[1, 2, 3.5], role="first input", requires_grad=True)
>>> y = Variable(data=[4, 5, 6], role="second input", requires_grad=False)
>>> result = F.sum([x, y])
>>> result.data
[5, 7, 9.5]
>>> result.role
'first input and second input'
>>> result.requires_grad
True

Modules

functional
modules
parameter