e3nn_mlx.o3.FullTensorProduct

class e3nn_mlx.o3.FullTensorProduct(irreps_in1, irreps_in2, filter_ir_out=None, *, irrep_normalization='component', path_normalization='element', compile_left_right=True, use_custom_kernel=True)

Raw-array compatible FullTensorProduct.

Parameters:

• irreps_in1 (Irreps | str)

• irreps_in2 (Irreps | str)

• filter_ir_out (list[Irrep | str] | None)

• irrep_normalization (str)

• path_normalization (str)

• compile_left_right (bool)

• use_custom_kernel (bool)

__init__(irreps_in1, irreps_in2, filter_ir_out=None, *, irrep_normalization='component', path_normalization='element', compile_left_right=True, use_custom_kernel=True)

Parameters:

• irreps_in1 (Irreps | str)

• irreps_in2 (Irreps | str)

• filter_ir_out (list[Irrep | str] | None)

• irrep_normalization (str)

• path_normalization (str)

• compile_left_right (bool)

• use_custom_kernel (bool)

Return type:

None

Methods

__init__(irreps_in1, irreps_in2[, ...])
apply(map_fn[, filter_fn])	Map all the parameters using the provided map_fn and immediately update the module with the mapped parameters.
apply_to_modules(apply_fn)	Apply a function to all the modules in this instance (including this instance).
children()	Return the direct descendants of this Module instance.
clear()
copy()
differentiable_arrays(left_array, right_array)	General MLX fallback for JVP and arbitrary transform nesting.
eval()	Set the model to evaluation mode.
filter_and_map(filter_fn[, map_fn, is_leaf_fn])	Recursively filter the contents of the module using filter_fn, namely only select keys and values where filter_fn returns true.
forward(left, right[, weight])
freeze(*[, recurse, keys, strict])	Freeze the Module's parameters or some of them.
fromkeys(iterable[, value])	Create a new dictionary with keys from iterable and values set to value.
get(key[, default])	Return the value for key if key is in the dictionary, else default.
is_module(value)
items()
keys()
leaf_modules()	Return the submodules that do not contain other modules.
load_weights(file_or_weights[, strict])	Update the model's weights from a .npz, a .safetensors file, or a list.
modules()	Return a list with all the modules in this instance.
named_modules()	Return a list with all the modules in this instance and their name with dot notation.
parameters()	Recursively return all the mlx.core.array members of this Module as a dict of dicts and lists.
pop(k[,d])	If the key is not found, return the default if given; otherwise, raise a KeyError.
popitem()	Remove and return a (key, value) pair as a 2-tuple.
right(right[, weight])
save_weights(file)	Save the model's weights to a file.
set_dtype(dtype[, predicate])	Set the dtype of the module's parameters.
setdefault(key[, default])	Insert key with a value of default if key is not in the dictionary.
train([mode])	Set the model in or out of training mode.
trainable_parameter_filter(module, key, value)
trainable_parameters()	Recursively return all the non frozen mlx.core.array members of this Module as a dict of dicts and lists.
unfreeze(*[, recurse, keys, strict])	Unfreeze the Module's parameters or some of them.
update(parameters[, strict])	Replace the parameters of this Module with the provided ones in the dict of dicts and lists.
update_modules(modules[, strict])	Replace the child modules of this Module instance with the provided ones in the dict of dicts and lists.
valid_child_filter(module, key, value)
valid_parameter_filter(module, key, value)
values()
weight_view_for_instruction(instruction_index)
weight_views([weight])

Attributes

output_mask
state	The module's state dictionary
training	Boolean indicating if the model is in training mode.