API Reference

Einsum Nodes

class feinsum.einsum.FusedEinsum(arg_shapes, value_to_dtype, access_descriptors, use_matrix, index_names)

A fused einsum expression.

shape

ndim

index_to_dim_length()

Return type

Map

get_subscripts()

Returns the subscripts used in the building the einsum from it.

Return type

str

get_arg_shape(name)

Returns the shape for argument named name.

Return type

Tuple[Union[int, integer, SizeParam], ...]

class feinsum.einsum.VeryLongAxis

Describes a dimension length which is to be assumed to be very large.

class feinsum.einsum.EinsumAxisAccess

Base class for axis access types in an einsum expression.

class feinsum.einsum.FreeAxis(output_index)

Records the axis of an einsum argument over which contraction is not performed.

output_index

Position of the corresponding index in the einsum’s output.

class feinsum.einsum.SummationAxis(index)

Records an index in an einsum expression over which reduction is performed. Sometimes also referred to as an axis with a corresponding “dummy index” in Ricci Calculus.

index

An integer which is unique to a reduction index of an einsum.

class feinsum.einsum.Argument

An abstract class denoting an argument to an einsum in ContractionSchedule. See ContractionSchedule.arguments.

class feinsum.einsum.EinsumOperand(ioperand)

An Argument representing the ioperand-th argument that was passed to the parent einsum whose ContractionSchedule is being specified.

class feinsum.einsum.IntermediateResult(name)

An Argument representing an intermediate result available during the current contraction.

class feinsum.einsum.ContractionSchedule(subscripts, result_names, arguments)

Records the schedule in which contractions are to be performed in an einsum as a series of einsums with the i-th einsum having subscript subscript[i] operating on arguments[i] and writing its result to result_names[i].

result_names

Names of the result generated by each step.

arguments

A tuple containing tuple of str for each contraction in the schedule.

nsteps

class feinsum.einsum.SizeParam(name)

Helper routines

feinsum.einsum.get_trivial_contraction_schedule(einsum)

Returns the ContractionSchedule for einsum scheduled as a single contraction.

Return type

ContractionSchedule

feinsum.einsum.get_opt_einsum_contraction_schedule(expr, **opt_einsum_kwargs)

Returns a ContractionSchedule as computed by opt_einsum.contract_path().

Parameters

opt_einsum_kwargs (Any) – kwargs to be passed to opt_einsum.contract_path().

Note

The following defaults are populated in opt_einsum_kwargs, if left unspecified:

• optimize="optimal"

• use_blas=False

Return type

ContractionSchedule

Routines for empirical evaluation

feinsum.measure.timeit(einsum, *, transform, cl_ctx, long_dim_length=100000, schedule=None)

Returns the runtime in seconds for executing einsum on OpenCL context cl_ctx.

Parameters

transform (TransformT) – The transformation to be applied to loopy.TranslationUnit lowered from einsum.

Return type

float

feinsum.measure.measure_giga_op_rate(expr, *, transform, cl_ctx, long_dim_length=100000, schedule=None)

Returns the arithmetic operations rate (in Giga Ops per second) by arithmetic operation’s result dtypes.

Return type

Map

feinsum.measure.stringify_comparison_vs_roofline(expr, *, schedule=None, transform, cl_ctx, long_dim_length=100000, ignore_unknown_device=False)

Returns the prettified comparison of expr transformed with transform wrt roofline. The roofline model assumes that kernel’s performance is limited by either the device’s global memory bandwidth or the device’s floating point units being saturated to their maximum throughput.

Parameters

ignore_unknown_device (bool) – If False raises an error if a roofline model is unknown for the device in cl_ctx. If True, no error is raised and the roofline performance is marked as “N/A” in the output.

Return type

str

Helpers to build feinsum.einsum.FusedEinsum

feinsum.make_einsum.einsum(subscripts, *operands, arg_names=None)

Returns a FusedEinsum with an interface similar to numpy.einsum().

Parameters

arg_names (Optional[Sequence[str]]) – An optional sequence of str. If not provided, defaults to the sequence: "arg_0", "arg_1", "arg_2", ....

Return type

FusedEinsum

feinsum.make_einsum.fused_einsum(subscripts, operand_shapes, use_matrix, dtypes=None, value_to_dtype=None)

Returns a FusedEinsum with an interface similar to numpy.einsum().

Parameters

• subscripts (str) – A str describing the Einstein summation as accepted by numpy.einsum().

• dtypes (Union[dtype[Any], None, Type[Any], _SupportsDType[dtype[Any]], str, Tuple[Any, int], Tuple[Any, Union[SupportsIndex, Sequence[SupportsIndex]]], List[Any], _DTypeDict, Tuple[Any, Any]]) – The dtype of all the value the operands use. Cannot be provide both value_to_dtype and dtypes.

• value_to_dtype (Optional[Mapping[str, Union[dtype[Any], None, Type[Any], _SupportsDType[dtype[Any]], str, Tuple[Any, int], Tuple[Any, Union[SupportsIndex, Sequence[SupportsIndex]]], List[Any], _DTypeDict, Tuple[Any, Any]]]]) – A mapping from the values the operands of the einsum depend on to their dtypes. Cannot be provide both value_to_dtype and dtypes.

• operand_shapes (Sequence[Any]) – A sequence of shapes of the operands of the Einstein Summation.

• use_matrix (Any) – A 2D numpy array-like object, where i,j-th entry corresponds to the frozenset of str of value names that the j-th operand of the i-th einsum accesses.

Return type

FusedEinsum

Interacting with code-transformation databases

feinsum.sql_utils.query(einsum, cl_ctx, *, database='/home/docs/checkouts/readthedocs.org/user_builds/feinsum/conda/latest/lib/python3.10/site-packages/feinsum/../../data/transform_archive_v3.sqlite', err_if_no_results=False)

Returns facts of previous recorded runs of einsum on cl_ctx.

Parameters

err_if_no_results (bool) – If True, raises a RuntimeError if no recorded runs corresponding to einsum are available in database. Defaults to False.

Return type

Tuple[QueryInfo, ...]

class feinsum.sql_utils.QueryInfo(transform_id, transform_params, runtime_in_sec, compiler_version, giga_op_info, _einsum)

Some helper loopy transformations relevant for Einstein Summations

loopy helpers. The aim is to keep this module as lean as possible by upstreaming the heavy lifting parts to loopy itself.

feinsum.loopy_utils.get_a_matched_einsum(t_unit, kernel_name=None, insn_match=None, argument_substitutions=None, long_dim_length=500)

Returns a tuple of the form (matched_einsum, subst_map) where, matched_einsum is the batched einsum having a memory access pattern similar to the instructions in insn_match of t_unit, and, subst_map is a mapping from the variables in t_unit to the entities (i.e. indices, arguments) of match_einsum.

Parameters

• t_unit (TranslationUnit) – The subject translation unit which is being matched against ref_einsum.

• insn_match (Optional[Any]) – A match expression as understood by loopy.match.parse_match() representing the subset of the kernel’s instructions raised to a batched einsum.

• argument_substitutions (Optional[FrozenSet[str]]) – A frozenset of substitutions that are to be matched as an Einstein-summation’s argument. If None, any substitution rule invocation is treated as accessing an Einstein Summation’s argument.

• long_dim_length (int) – Axis length above which can be assumed to be long enough to be invariant of having any noticeable effects on the kernel’s performance.

Note

See xxx (TODO) for a description of a translation unit with a einsum-like memory access pattern.

Return type

Tuple[FusedEinsum, frozenbidict[str, str]]

feinsum.loopy_utils.match_t_unit_to_einsum(t_unit, einsum, *, kernel_name=None, insn_match=None, argument_substitutions=None, long_dim_length=500)

Returns a mapping from the entities of einsum to the variables of the corresponding matched einsum in t_unit. See get_a_matched_einsum() for a subset of grammar of loopy kernels to be a matched as a batched einsum.

Return type

Mapping[str, str]

feinsum.loopy_utils.get_call_ids(expr)

Returns the identifiers of the invoked functions in expr.

Return type

FrozenSet[str]

Lowering to loopy

feinsum.codegen.loopy.generate_loopy(einsum, schedule=None)

Returns a loopy.TranslationUnit with the reductions scheduled by contract_path.

Parameters

schedule (Optional[ContractionSchedule]) – An optional instance of ContractionSchedule. Defaults to the trivial contraction schedule if not provided.

Return type

TranslationUnit

feinsum.codegen.loopy.generate_loopy_with_opt_einsum_schedule(expr, **opt_einsum_kwargs)

Returns a loopy.TranslationUnit with the ContractionSchedule specified via opt_einsum.contract_path().

Parameters

opt_einsum_kwargs (Any) – kwargs to be passed to get_opt_einsum_contraction_schedule().

Return type

TranslationUnit

Canonicalization

feinsum.canonicalization.canonicalize_einsum(einsum)

Returns a canonicalized form of einsum.

Note

• Refer to (TODO PAPER) for a definition of isomorphism among fused einsums.

Return type

FusedEinsum

Miscellaneous Utilities

feinsum.utils.has_similar_subscript(einsum, subscript)

Returns True only if einsum’s expression being applied to its operands is equivalently represented by subscript.

>>> import feinsum as f
>>> ensm = f.einsum("ij,j->i",
...                 f.array((10, 4), "float64"),
...                 f.array(4, "float64"))
>>> has_similar_subscript(ensm, "ij,j->i")
True
>>> has_similar_subscript(ensm, "ik,k->i")
True
>>> has_similar_subscript(ensm, "ik,kj->ij")
False
>>> ensm = f.einsum("ik,kj->ij",
...                 f.array((10, 4), "float64"),
...                 f.array((4, 100), "float64"))
>>> has_similar_subscript(ensm, "ij,j->i")
False
>>> has_similar_subscript(ensm, "ik,kj->ij")
True
>>> has_similar_subscript(ensm, "pq,qr->pr")
True

Return type

bool

feinsum.utils.is_any_redn_dim_parametric(einsum)

>>> import feinsum as f
>>> ensm = f.einsum("ij,j->i",
...                 f.array((10, 4), "float64"),
...                 f.array(4, "float64"))
>>> is_any_redn_dim_parametric(ensm)
False
>>> ensm = f.einsum("ij,jk->i",
...                 f.array((10, 4), "float64"),
...                 f.array((4, 20), "float64"))
>>> is_any_redn_dim_parametric(ensm)
False
>>> ensm = f.einsum("i,j->i",
...                 f.array(125, "float64"),
...                 f.array(np.inf, "float64"))
>>> is_any_redn_dim_parametric(ensm)
True

Return type

bool

class feinsum.typing.TransformT(*args, **kwargs)

Indices and tables

• Index

• Module Index