product

tangermeme.product.apply_product(func: Callable[[...], Any], model: Module, X: Tensor, args: Sequence[Tensor], batch_size: int = 32, device: str | device | None = None, additional_func_kwargs: dict | None = None, verbose: bool = False, **kwargs: Any) → Tensor | list[Tensor]

Apply a function on the cartesian product between X and each args.

This function will take the provided function and apply it in a batched manner across the cartesian product of X and each of the arguments provided in args. Because this is a cartesian product, the number of examples that need to be processed will quickly grow with respect to the number of arguments being passed in. Each of the tensors in args must be one input to model, in the order that they are specified by the forward function.

This function can accept in any other function – be it predictions, attributions, or marginalizations. If the provided function itself has parameters that need to be specified, you can provide them directly to this function in the order that they appear in the provided function.

Parameters

func: function

A function, likely implemented in tangermeme, to apply in a batched manner across the product of examples.

model: torch.nn.Module

The PyTorch model to use to make predictions.

X: torch.tensor, shape=(-1, len(alphabet), length)

A one-hot encoded set of sequences to make predictions for.

args: tuple or list

A set of additional arguments to pass into the model. Each element in args should be one tensor that is input to the model. The elements do not need to be the same size as each other as a product will be constructed over all of them, as well as with X. If you only want to use one value for an argument across all function applications, pass it in as a length-1 tensor.

batch_size: int, optional

The number of examples to make predictions for at a time. Default is 32.

device: str or torch.device or None, optional

The device to move the model and batches to when making predictions. If None, use CUDA when available and fall back to CPU otherwise. The model’s original device and training mode are restored after the call. Default is None.

additional_func_kwargs: dict or None, optional

Additional named arguments to pass into the function when it is called. This is provided as an alternate path to route arguments into the function in case they overlap, name-wise, with those in this function, or if you want to be absolutely sure that the arguments are making their way into the function. The dict is not modified in place. Note that overlap between keys of additional_func_kwargs and **kwargs raises TypeError: multiple values for keyword argument from Python’s call-site dict-unpacking; only collisions with this function’s named parameters are resolved by routing through additional_func_kwargs. Default is None.

verbose: bool, optional

Whether to display a progress bar as examples are processed. Default is False.

kwargs: optional

Additional named arguments that will get passed into the function when it is called. Default is no arguments are passed in.

Returns

y: torch.Tensor or list/tuple of torch.Tensors

The output from the model for each input example. The precise format is determined by the model. If the model outputs a single tensor, y is a single tensor concatenated across all batches. If the model outputs multiple tensors, y is a list of tensors which are each concatenated across all batches.