Skip to content

Counted Operations

Operations that consume the FLOP budget. Every call deducts its analytical cost from the active budget before execution.

Cost rules of thumb:

  • Unary (exp, log, sqrt, ...): 1 FLOP per output element
  • Binary (add, multiply, maximum, ...): 1 FLOP per output element (after broadcasting)
  • Reductions (sum, mean, max, ...): 1 FLOP per input element
  • Einsum: product_of_all_index_dims x op_factor (op_factor=2 for inner products)

See FLOP Counting Model for full details.

API Reference

mechestim._einsum.einsum(subscripts, *operands, optimize='auto', symmetric_axes=None, **kwargs)

Evaluate Einstein summation with FLOP counting and optional path optimization.

Wraps numpy.einsum with analytical FLOP cost computation and optional symmetry savings. If any input is a SymmetricTensor, the cost is automatically reduced. If symmetric_axes is provided and the output passes validation, a SymmetricTensor is returned.

All contractions go through opt_einsum's contract_path to find an optimal pairwise decomposition. The FLOP cost uses opt_einsum's cost model which includes op_factor (multiply-add = 2 FLOPs for inner products).

Parameters:

Name Type Description Default
subscripts str

Einstein summation subscript string (e.g., 'ij,jk->ik').

 required
*operands ndarray

Input arrays. SymmetricTensor inputs are detected automatically for cost savings.

 ()
optimize str, bool, or list of tuple

Contraction path strategy. Default 'auto'.

  • 'auto', 'greedy', 'optimal', 'dp', etc.: Use the named algorithm to find the best path.
  • A list of int-tuples (e.g. [(1, 2), (0, 1)]): use this explicit contraction path. Obtain one from me.einsum_path() or construct manually. Each tuple names the operand positions to contract at that step; the result is appended to the end.
  • False: treated as 'auto'.
 'auto'
symmetric_axes list of tuple of int

Output dimension symmetry groups. Declares that the result is symmetric in the given axes and wraps it as a SymmetricTensor. For example, [(0, 1)] means the output satisfies result[i,j,...] == result[j,i,...]. This does NOT declare input symmetry — use me.as_symmetric() for that.

 None

Returns:

Type Description
ndarray or SymmetricTensor

The result of the einsum.

Raises:

Type Description
BudgetExhaustedError

If the operation would exceed the FLOP budget.
NoBudgetContextError

If called outside a BudgetContext.
SymmetryError

If symmetric_axes is provided but the result is not symmetric.
Source code in src/mechestim/_einsum.py 
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def einsum(
    subscripts: str,
    *operands: _np.ndarray,
    optimize: str | bool | list = "auto",
    symmetric_axes: list[tuple[int, ...]] | None = None,
    **kwargs,
) -> _np.ndarray:
    """Evaluate Einstein summation with FLOP counting and optional path optimization.

    Wraps ``numpy.einsum`` with analytical FLOP cost computation and
    optional symmetry savings. If any input is a ``SymmetricTensor``,
    the cost is automatically reduced. If ``symmetric_axes`` is provided
    and the output passes validation, a ``SymmetricTensor`` is returned.

    All contractions go through opt_einsum's ``contract_path`` to find an
    optimal pairwise decomposition. The FLOP cost uses opt_einsum's cost
    model which includes ``op_factor`` (multiply-add = 2 FLOPs for inner
    products).

    Parameters
    ----------
    subscripts : str
        Einstein summation subscript string (e.g., ``'ij,jk->ik'``).
    *operands : numpy.ndarray
        Input arrays. ``SymmetricTensor`` inputs are detected automatically
        for cost savings.
    optimize : str, bool, or list of tuple, optional
        Contraction path strategy. Default ``'auto'``.

        - ``'auto'``, ``'greedy'``, ``'optimal'``, ``'dp'``, etc.:
          Use the named algorithm to find the best path.
        - A list of int-tuples (e.g. ``[(1, 2), (0, 1)]``): use this
          explicit contraction path. Obtain one from ``me.einsum_path()``
          or construct manually. Each tuple names the operand positions
          to contract at that step; the result is appended to the end.
        - ``False``: treated as ``'auto'``.
    symmetric_axes : list of tuple of int, optional
        **Output** dimension symmetry groups. Declares that the result
        is symmetric in the given axes and wraps it as a
        ``SymmetricTensor``. For example, ``[(0, 1)]`` means the output
        satisfies ``result[i,j,...] == result[j,i,...]``. This does NOT
        declare input symmetry — use ``me.as_symmetric()`` for that.

    Returns
    -------
    numpy.ndarray or SymmetricTensor
        The result of the einsum.

    Raises
    ------
    BudgetExhaustedError
        If the operation would exceed the FLOP budget.
    NoBudgetContextError
        If called outside a ``BudgetContext``.
    SymmetryError
        If ``symmetric_axes`` is provided but the result is not symmetric.
    """
    budget = require_budget()
    shapes = [op.shape for op in operands]

    # Extract symmetry info from SymmetricTensor inputs
    operand_symmetries = [
        op.symmetry_info if isinstance(op, SymmetricTensor) else None for op in operands
    ]

    # Convert SymmetryInfo -> IndexSymmetry for the path optimizer
    input_parts = subscripts.split("->")[0].split(",")
    index_symmetries = [
        _symmetry_info_to_index_symmetry(s, chars)
        for s, chars in zip(operand_symmetries, input_parts)
    ]
    has_symmetry = any(s is not None for s in index_symmetries)

    from mechestim._opt_einsum import contract_path as _contract_path

    path, path_info = _contract_path(
        subscripts,
        *shapes,
        shapes=True,
        optimize=optimize if optimize is not False else "auto",
        input_symmetries=index_symmetries if has_symmetry else None,
    )

    budget.deduct(
        "einsum",
        flop_cost=path_info.optimized_cost,
        subscripts=subscripts,
        shapes=tuple(shapes),
    )

    # Execute pairwise steps
    result = _execute_pairwise(path_info, list(operands))

    # Handle output symmetry wrapping
    if symmetric_axes and isinstance(result, _np.ndarray) and result.ndim >= 2:
        validate_symmetry(result, symmetric_axes)
        result = SymmetricTensor(result, symmetric_axes=symmetric_axes)

    check_nan_inf(result, "einsum")
    return result

mechestim._einsum.einsum_path(subscripts, *operands, optimize='auto')

Compute the optimal contraction path without executing.

Returns (path, PathInfo) with zero budget cost. The returned path can be passed back to me.einsum(..., optimize=path) to execute with that exact contraction order.

Parameters:

Name Type Description Default
subscripts str

Einstein summation subscript string.

 required
*operands ndarray

Input arrays.

 ()
optimize str, bool, or list of tuple

Path optimization strategy. Default 'auto'.

 'auto'

Returns:

Name Type Description
path list of tuple of int

The contraction path. Pass to me.einsum(..., optimize=path).
info PathInfo

Diagnostics including per-step costs and symmetry savings.
Source code in src/mechestim/_einsum.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def einsum_path(subscripts: str, *operands, optimize: str | bool | list = "auto"):
    """Compute the optimal contraction path without executing.

    Returns ``(path, PathInfo)`` with zero budget cost. The returned
    ``path`` can be passed back to ``me.einsum(..., optimize=path)``
    to execute with that exact contraction order.

    Parameters
    ----------
    subscripts : str
        Einstein summation subscript string.
    *operands : numpy.ndarray
        Input arrays.
    optimize : str, bool, or list of tuple, optional
        Path optimization strategy. Default ``'auto'``.

    Returns
    -------
    path : list of tuple of int
        The contraction path. Pass to ``me.einsum(..., optimize=path)``.
    info : PathInfo
        Diagnostics including per-step costs and symmetry savings.
    """
    shapes = [op.shape for op in operands]
    operand_symmetries = [
        op.symmetry_info if isinstance(op, SymmetricTensor) else None for op in operands
    ]
    input_parts = subscripts.split("->")[0].split(",")
    index_symmetries = [
        _symmetry_info_to_index_symmetry(s, chars)
        for s, chars in zip(operand_symmetries, input_parts)
    ]
    has_symmetry = any(s is not None for s in index_symmetries)

    from mechestim._opt_einsum import contract_path as _contract_path

    path, path_info = _contract_path(
        subscripts,
        *shapes,
        shapes=True,
        optimize=optimize if optimize is not False else "auto",
        input_symmetries=index_symmetries if has_symmetry else None,
    )
    return list(path), path_info

mechestim._pointwise

Counted pointwise operations and reductions for mechestim.

around(a, decimals=0, out=None)

Counted version of np.around. Cost = numel(output) FLOPs.

Source code in src/mechestim/_pointwise.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def around(a, decimals=0, out=None):
    """Counted version of np.around. Cost = numel(output) FLOPs."""
    budget = require_budget()
    a_is_scalar = not isinstance(a, _np.ndarray) and _np.ndim(a) == 0
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    sym_info = a.symmetry_info if isinstance(a, SymmetricTensor) else None
    cost = pointwise_cost(a.shape, symmetry_info=sym_info)
    budget.deduct("around", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    result = _np.around(a, decimals=decimals, out=out)
    check_nan_inf(result, "around")
    if sym_info is not None:
        result = SymmetricTensor(result, symmetric_axes=sym_info.symmetric_axes)
    if (
        a_is_scalar
        and out is None
        and isinstance(result, _np.ndarray)
        and result.ndim == 0
    ):
        return result.item()
    return result

clip(a, a_min=None, a_max=None, out=None, **kwargs)

Counted version of np.clip. Cost = numel(input) or unique_elements if symmetric.

Source code in src/mechestim/_pointwise.py 
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
def clip(a, a_min=None, a_max=None, out=None, **kwargs):
    """Counted version of np.clip. Cost = numel(input) or unique_elements if symmetric."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    sym_info = a.symmetry_info if isinstance(a, SymmetricTensor) else None
    cost = pointwise_cost(a.shape, symmetry_info=sym_info)
    budget.deduct("clip", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    # numpy forbids min=/max= kwargs when a_min/a_max positional args are given;
    # handle the case where caller uses min=/max= keyword style
    if "min" in kwargs or "max" in kwargs:
        if a_min is None and "min" in kwargs:
            a_min = kwargs.pop("min")
        if a_max is None and "max" in kwargs:
            a_max = kwargs.pop("max")
    result = _np.clip(a, a_min, a_max, out=out, **kwargs)
    if a.dtype.kind in ("f", "c"):
        check_nan_inf(result, "clip")
    if sym_info is not None:
        result = SymmetricTensor(result, symmetric_axes=sym_info.symmetric_axes)
    return result

convolve(a, v, mode='full')

Counted version of np.convolve.

Source code in src/mechestim/_pointwise.py 
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
def convolve(a, v, mode="full"):
    """Counted version of np.convolve."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(v, _np.ndarray):
        v = _np.asarray(v)
    result = _np.convolve(a, v, mode=mode)
    budget.deduct(
        "convolve",
        flop_cost=_builtins.max(a.size * v.size, 1),
        subscripts=None,
        shapes=(a.shape, v.shape),
    )
    return result

corrcoef(x, y=None, **kwargs)

Counted version of np.corrcoef.

Source code in src/mechestim/_pointwise.py 
818
819
820
821
822
823
824
825
826
827
def corrcoef(x, y=None, **kwargs):
    """Counted version of np.corrcoef."""
    budget = require_budget()
    if not isinstance(x, _np.ndarray):
        x = _np.asarray(x)
    cost = x.size * x.size if y is None else x.size * _np.asarray(y).size
    budget.deduct(
        "corrcoef", flop_cost=_builtins.max(cost, 1), subscripts=None, shapes=(x.shape,)
    )
    return _np.corrcoef(x, y=y, **kwargs)

correlate(a, v, mode='valid')

Counted version of np.correlate.

Source code in src/mechestim/_pointwise.py 
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
def correlate(a, v, mode="valid"):
    """Counted version of np.correlate."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(v, _np.ndarray):
        v = _np.asarray(v)
    result = _np.correlate(a, v, mode=mode)
    budget.deduct(
        "correlate",
        flop_cost=_builtins.max(a.size * v.size, 1),
        subscripts=None,
        shapes=(a.shape, v.shape),
    )
    return result

cov(m, y=None, **kwargs)

Counted version of np.cov.

Source code in src/mechestim/_pointwise.py 
833
834
835
836
837
838
839
840
841
842
def cov(m, y=None, **kwargs):
    """Counted version of np.cov."""
    budget = require_budget()
    if not isinstance(m, _np.ndarray):
        m = _np.asarray(m)
    cost = m.size * m.size if y is None else m.size * _np.asarray(y).size
    budget.deduct(
        "cov", flop_cost=_builtins.max(cost, 1), subscripts=None, shapes=(m.shape,)
    )
    return _np.cov(m, y=y, **kwargs)

cross(a, b, **kwargs)

Counted version of np.cross.

Source code in src/mechestim/_pointwise.py 
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
def cross(a, b, **kwargs):
    """Counted version of np.cross."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    result = _np.cross(a, b, **kwargs)
    r = _np.asarray(result)
    budget.deduct(
        "cross",
        flop_cost=_builtins.max(r.size * 3, 1),
        subscripts=None,
        shapes=(a.shape, b.shape),
    )
    return result

diff(a, n=1, axis=-1, **kwargs)

Counted version of np.diff.

Source code in src/mechestim/_pointwise.py 
730
731
732
733
734
735
736
737
738
739
740
741
742
def diff(a, n=1, axis=-1, **kwargs):
    """Counted version of np.diff."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    result = _np.diff(a, n=n, axis=axis, **kwargs)
    budget.deduct(
        "diff",
        flop_cost=_builtins.max(result.size, 1),
        subscripts=None,
        shapes=(a.shape,),
    )
    return result

dot(a, b)

Counted version of np.dot.

Source code in src/mechestim/_pointwise.py 
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
def dot(a, b):
    """Counted version of np.dot."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    # Extract symmetry info for cost calculation
    operand_symmetries = [
        a.symmetry_info if isinstance(a, SymmetricTensor) else None,
        b.symmetry_info if isinstance(b, SymmetricTensor) else None,
    ]
    has_sym = _builtins.any(s is not None for s in operand_symmetries)
    if a.ndim == 2 and b.ndim == 2:
        cost = einsum_cost(
            "ij,jk->ik",
            shapes=[a.shape, b.shape],
            operand_symmetries=operand_symmetries if has_sym else None,
        )
    elif a.ndim == 1 and b.ndim == 1:
        cost = einsum_cost(
            "i,i->",
            shapes=[a.shape, b.shape],
            operand_symmetries=operand_symmetries if has_sym else None,
        )
    else:
        cost = a.size * b.size
    budget.deduct("dot", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape))
    result = _np.dot(a, b)
    check_nan_inf(result, "dot")
    return result

ediff1d(ary, **kwargs)

Counted version of np.ediff1d.

Source code in src/mechestim/_pointwise.py 
760
761
762
763
764
765
766
767
768
769
770
771
772
def ediff1d(ary, **kwargs):
    """Counted version of np.ediff1d."""
    budget = require_budget()
    if not isinstance(ary, _np.ndarray):
        ary = _np.asarray(ary)
    result = _np.ediff1d(ary, **kwargs)
    budget.deduct(
        "ediff1d",
        flop_cost=_builtins.max(result.size, 1),
        subscripts=None,
        shapes=(ary.shape,),
    )
    return result

gradient(f, *varargs, **kwargs)

Counted version of np.gradient.

Source code in src/mechestim/_pointwise.py 
748
749
750
751
752
753
754
def gradient(f, *varargs, **kwargs):
    """Counted version of np.gradient."""
    budget = require_budget()
    if not isinstance(f, _np.ndarray):
        f = _np.asarray(f)
    budget.deduct("gradient", flop_cost=f.size, subscripts=None, shapes=(f.shape,))
    return _np.gradient(f, *varargs, **kwargs)

inner(a, b)

Counted version of np.inner.

Source code in src/mechestim/_pointwise.py 
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
def inner(a, b):
    """Counted version of np.inner."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    cost = (
        a.size
        if (a.ndim <= 1 and b.ndim <= 1)
        else a.size * (b.shape[-1] if b.ndim > 1 else 1)
    )
    budget.deduct("inner", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape))
    result = _np.inner(a, b)
    return result

interp(x, xp, fp, **kwargs)

Counted version of np.interp.

Source code in src/mechestim/_pointwise.py 
872
873
874
875
876
877
878
879
880
def interp(x, xp, fp, **kwargs):
    """Counted version of np.interp."""
    budget = require_budget()
    if not isinstance(x, _np.ndarray):
        x = _np.asarray(x)
    budget.deduct(
        "interp", flop_cost=_builtins.max(x.size, 1), subscripts=None, shapes=(x.shape,)
    )
    return _np.interp(x, xp, fp, **kwargs)

isclose(a, b, **kwargs)

Counted version of np.isclose. Cost = numel(output) FLOPs.

Source code in src/mechestim/_pointwise.py 
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
def isclose(a, b, **kwargs):
    """Counted version of np.isclose. Cost = numel(output) FLOPs."""
    budget = require_budget()
    a_is_scalar = not isinstance(a, _np.ndarray) and _np.ndim(a) == 0
    b_is_scalar = not isinstance(b, _np.ndarray) and _np.ndim(b) == 0
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    output_shape = _np.broadcast_shapes(a.shape, b.shape)
    cost = pointwise_cost(output_shape)
    budget.deduct("isclose", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape))
    result = _np.isclose(a, b, **kwargs)
    if (
        a_is_scalar
        and b_is_scalar
        and isinstance(result, _np.ndarray)
        and result.ndim == 0
    ):
        return result.item()
    return result

kron(a, b)

Counted version of np.kron.

Source code in src/mechestim/_pointwise.py 
692
693
694
695
696
697
698
699
700
701
702
703
def kron(a, b):
    """Counted version of np.kron."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    result = _np.kron(a, b)
    budget.deduct(
        "kron", flop_cost=result.size, subscripts=None, shapes=(a.shape, b.shape)
    )
    return result

matmul(a, b)

Counted version of np.matmul.

Source code in src/mechestim/_pointwise.py 
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
def matmul(a, b):
    """Counted version of np.matmul."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    # Extract symmetry info for cost calculation
    operand_symmetries = [
        a.symmetry_info if isinstance(a, SymmetricTensor) else None,
        b.symmetry_info if isinstance(b, SymmetricTensor) else None,
    ]
    has_sym = _builtins.any(s is not None for s in operand_symmetries)
    if a.ndim == 2 and b.ndim == 2:
        cost = einsum_cost(
            "ij,jk->ik",
            shapes=[a.shape, b.shape],
            operand_symmetries=operand_symmetries if has_sym else None,
        )
    elif a.ndim == 1 and b.ndim == 1:
        cost = einsum_cost(
            "i,i->",
            shapes=[a.shape, b.shape],
            operand_symmetries=operand_symmetries if has_sym else None,
        )
    else:
        cost = a.size * b.size
    budget.deduct("matmul", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape))
    result = _np.matmul(a, b)
    check_nan_inf(result, "matmul")
    return result

outer(a, b, out=None)

Counted version of np.outer.

Source code in src/mechestim/_pointwise.py 
635
636
637
638
639
640
641
642
643
644
def outer(a, b, out=None):
    """Counted version of np.outer."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    cost = a.size * b.size
    budget.deduct("outer", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape))
    return _np.outer(a, b, out=out)

ptp(a, axis=None, **kwargs)

Peak-to-peak range. Cost = numel(input) FLOPs.

Source code in src/mechestim/_pointwise.py 
521
522
523
524
525
526
527
528
def ptp(a, axis=None, **kwargs):
    """Peak-to-peak range. Cost = numel(input) FLOPs."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    cost = reduction_cost(a.shape, axis)
    budget.deduct("ptp", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    return _np.max(a, axis=axis, **kwargs) - _np.min(a, axis=axis, **kwargs)

round(a, decimals=0, out=None)

Counted version of np.round. Cost = numel(output) FLOPs.

Source code in src/mechestim/_pointwise.py 
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def round(a, decimals=0, out=None):
    """Counted version of np.round. Cost = numel(output) FLOPs."""
    budget = require_budget()
    a_is_scalar = not isinstance(a, _np.ndarray) and _np.ndim(a) == 0
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    sym_info = a.symmetry_info if isinstance(a, SymmetricTensor) else None
    cost = pointwise_cost(a.shape, symmetry_info=sym_info)
    budget.deduct("round", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    result = _np.round(a, decimals=decimals, out=out)
    check_nan_inf(result, "round")
    if sym_info is not None:
        result = SymmetricTensor(result, symmetric_axes=sym_info.symmetric_axes)
    if (
        a_is_scalar
        and out is None
        and isinstance(result, _np.ndarray)
        and result.ndim == 0
    ):
        return result.item()
    return result

tensordot(a, b, axes=2)

Counted version of np.tensordot.

Source code in src/mechestim/_pointwise.py 
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
def tensordot(a, b, axes=2):
    """Counted version of np.tensordot."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    result = _np.tensordot(a, b, axes=axes)
    if isinstance(axes, int):
        contracted = 1
        for i in range(axes):
            contracted *= a.shape[a.ndim - axes + i]
        cost = _builtins.max(result.size * contracted, 1)
    else:
        contracted = 1
        for i in axes[0]:
            contracted *= a.shape[i]
        cost = _builtins.max(result.size * contracted, 1)
    budget.deduct(
        "tensordot", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape)
    )
    return result

trapezoid(y, x=None, dx=1.0, axis=-1)

Counted version of np.trapezoid.

Source code in src/mechestim/_pointwise.py 
848
849
850
851
852
853
854
def trapezoid(y, x=None, dx=1.0, axis=-1):
    """Counted version of np.trapezoid."""
    budget = require_budget()
    if not isinstance(y, _np.ndarray):
        y = _np.asarray(y)
    budget.deduct("trapezoid", flop_cost=y.size, subscripts=None, shapes=(y.shape,))
    return _np.trapezoid(y, x=x, dx=dx, axis=axis)

trapz(y, x=None, dx=1.0, axis=-1)

Counted version of np.trapz (deprecated alias for trapezoid).

Source code in src/mechestim/_pointwise.py 
860
861
862
863
864
865
866
def trapz(y, x=None, dx=1.0, axis=-1):
    """Counted version of np.trapz (deprecated alias for trapezoid)."""
    budget = require_budget()
    if not isinstance(y, _np.ndarray):
        y = _np.asarray(y)
    budget.deduct("trapz", flop_cost=y.size, subscripts=None, shapes=(y.shape,))
    return _np.trapz(y, x=x, dx=dx, axis=axis)

vdot(a, b)

Counted version of np.vdot.

Source code in src/mechestim/_pointwise.py 
677
678
679
680
681
682
683
684
685
686
def vdot(a, b):
    """Counted version of np.vdot."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if not isinstance(b, _np.ndarray):
        b = _np.asarray(b)
    cost = a.size
    budget.deduct("vdot", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape))
    return _np.vdot(a, b)

mechestim._polynomial

Counted polynomial operations for mechestim.

poly(seq_of_zeros)

Return polynomial coefficients from roots. Wraps numpy.poly.

Source code in src/mechestim/_polynomial.py 
194
195
196
197
198
199
200
201
202
203
204
205
def poly(seq_of_zeros):
    """Return polynomial coefficients from roots. Wraps ``numpy.poly``."""
    budget = require_budget()
    seq = _np.asarray(seq_of_zeros)
    # If 2D (square matrix), n = shape[0]; if 1D, n = len(seq)
    if seq.ndim == 2:
        n = seq.shape[0]
    else:
        n = len(seq)
    cost = poly_cost(n)
    budget.deduct("poly", flop_cost=cost, subscripts=None, shapes=(seq.shape,))
    return _np.poly(seq_of_zeros)

poly_cost(n)

Cost for poly: \(n^2\) FLOPs.

Source code in src/mechestim/_polynomial.py 
55
56
57
def poly_cost(n: int) -> int:
    """Cost for poly: $n^2$ FLOPs."""
    return max(n * n, 1)

polyadd(a1, a2)

Add two polynomials. Wraps numpy.polyadd.

Source code in src/mechestim/_polynomial.py 
87
88
89
90
91
92
93
94
95
96
97
98
def polyadd(a1, a2):
    """Add two polynomials. Wraps ``numpy.polyadd``."""
    budget = require_budget()
    a1 = _np.asarray(a1)
    a2 = _np.asarray(a2)
    n1 = len(a1)
    n2 = len(a2)
    cost = polyadd_cost(n1, n2)
    budget.deduct(
        "polyadd", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape)
    )
    return _np.polyadd(a1, a2)

polyadd_cost(n1, n2)

Cost for polyadd: max(n1, n2) FLOPs.

Source code in src/mechestim/_polynomial.py 
20
21
22
def polyadd_cost(n1: int, n2: int) -> int:
    """Cost for polyadd: max(n1, n2) FLOPs."""
    return max(n1, n2, 1)

polyder(p, m=1)

Differentiate a polynomial. Wraps numpy.polyder.

Source code in src/mechestim/_polynomial.py 
121
122
123
124
125
126
127
128
def polyder(p, m=1):
    """Differentiate a polynomial. Wraps ``numpy.polyder``."""
    budget = require_budget()
    p = _np.asarray(p)
    n = len(p)
    cost = polyder_cost(n)
    budget.deduct("polyder", flop_cost=cost, subscripts=None, shapes=(p.shape,))
    return _np.polyder(p, m=m)

polyder_cost(n)

Cost for polyder: n FLOPs (n = len of coeffs).

Source code in src/mechestim/_polynomial.py 
30
31
32
def polyder_cost(n: int) -> int:
    """Cost for polyder: n FLOPs (n = len of coeffs)."""
    return max(n, 1)

polydiv(u, v)

Divide one polynomial by another. Wraps numpy.polydiv.

Source code in src/mechestim/_polynomial.py 
166
167
168
169
170
171
172
173
174
175
def polydiv(u, v):
    """Divide one polynomial by another. Wraps ``numpy.polydiv``."""
    budget = require_budget()
    u = _np.asarray(u)
    v = _np.asarray(v)
    n1 = len(u)
    n2 = len(v)
    cost = polydiv_cost(n1, n2)
    budget.deduct("polydiv", flop_cost=cost, subscripts=None, shapes=(u.shape, v.shape))
    return _np.polydiv(u, v)

polydiv_cost(n1, n2)

Cost for polydiv: n1 * n2 FLOPs.

Source code in src/mechestim/_polynomial.py 
45
46
47
def polydiv_cost(n1: int, n2: int) -> int:
    """Cost for polydiv: n1 * n2 FLOPs."""
    return max(n1 * n2, 1)

polyfit(x, y, deg, **kwargs)

Least-squares polynomial fit. Wraps numpy.polyfit.

Source code in src/mechestim/_polynomial.py 
181
182
183
184
185
186
187
188
def polyfit(x, y, deg, **kwargs):
    """Least-squares polynomial fit. Wraps ``numpy.polyfit``."""
    budget = require_budget()
    x = _np.asarray(x)
    m = len(x)
    cost = polyfit_cost(m, deg)
    budget.deduct("polyfit", flop_cost=cost, subscripts=None, shapes=(x.shape,))
    return _np.polyfit(x, y, deg, **kwargs)

polyfit_cost(m, deg)

Cost for polyfit: 2 * m * (deg+1)^2 FLOPs.

Source code in src/mechestim/_polynomial.py 
50
51
52
def polyfit_cost(m: int, deg: int) -> int:
    """Cost for polyfit: 2 * m * (deg+1)^2 FLOPs."""
    return max(2 * m * (deg + 1) ** 2, 1)

polyint(p, m=1, k=None)

Integrate a polynomial. Wraps numpy.polyint.

Source code in src/mechestim/_polynomial.py 
134
135
136
137
138
139
140
141
142
143
def polyint(p, m=1, k=None):
    """Integrate a polynomial. Wraps ``numpy.polyint``."""
    budget = require_budget()
    p = _np.asarray(p)
    n = len(p)
    cost = polyint_cost(n)
    budget.deduct("polyint", flop_cost=cost, subscripts=None, shapes=(p.shape,))
    if k is None:
        return _np.polyint(p, m=m)
    return _np.polyint(p, m=m, k=k)

polyint_cost(n)

Cost for polyint: n FLOPs (n = len of coeffs).

Source code in src/mechestim/_polynomial.py 
35
36
37
def polyint_cost(n: int) -> int:
    """Cost for polyint: n FLOPs (n = len of coeffs)."""
    return max(n, 1)

polymul(a1, a2)

Multiply polynomials. Wraps numpy.polymul.

Source code in src/mechestim/_polynomial.py 
149
150
151
152
153
154
155
156
157
158
159
160
def polymul(a1, a2):
    """Multiply polynomials. Wraps ``numpy.polymul``."""
    budget = require_budget()
    a1 = _np.asarray(a1)
    a2 = _np.asarray(a2)
    n1 = len(a1)
    n2 = len(a2)
    cost = polymul_cost(n1, n2)
    budget.deduct(
        "polymul", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape)
    )
    return _np.polymul(a1, a2)

polymul_cost(n1, n2)

Cost for polymul: n1 * n2 FLOPs.

Source code in src/mechestim/_polynomial.py 
40
41
42
def polymul_cost(n1: int, n2: int) -> int:
    """Cost for polymul: n1 * n2 FLOPs."""
    return max(n1 * n2, 1)

polysub(a1, a2)

Subtract two polynomials. Wraps numpy.polysub.

Source code in src/mechestim/_polynomial.py 
104
105
106
107
108
109
110
111
112
113
114
115
def polysub(a1, a2):
    """Subtract two polynomials. Wraps ``numpy.polysub``."""
    budget = require_budget()
    a1 = _np.asarray(a1)
    a2 = _np.asarray(a2)
    n1 = len(a1)
    n2 = len(a2)
    cost = polysub_cost(n1, n2)
    budget.deduct(
        "polysub", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape)
    )
    return _np.polysub(a1, a2)

polysub_cost(n1, n2)

Cost for polysub: max(n1, n2) FLOPs.

Source code in src/mechestim/_polynomial.py 
25
26
27
def polysub_cost(n1: int, n2: int) -> int:
    """Cost for polysub: max(n1, n2) FLOPs."""
    return max(n1, n2, 1)

polyval(p, x)

Evaluate a polynomial at given points. Wraps numpy.polyval.

Source code in src/mechestim/_polynomial.py 
70
71
72
73
74
75
76
77
78
79
def polyval(p, x):
    """Evaluate a polynomial at given points. Wraps ``numpy.polyval``."""
    budget = require_budget()
    p = _np.asarray(p)
    x = _np.asarray(x)
    deg = len(p) - 1
    m = x.size
    cost = polyval_cost(deg, m)
    budget.deduct("polyval", flop_cost=cost, subscripts=None, shapes=(p.shape, x.shape))
    return _np.polyval(p, x)

polyval_cost(deg, m)

Cost for polyval: Horner's method = 2 * m * deg FLOPs.

Source code in src/mechestim/_polynomial.py 
15
16
17
def polyval_cost(deg: int, m: int) -> int:
    """Cost for polyval: Horner's method = 2 * m * deg FLOPs."""
    return max(2 * m * deg, 1)

roots(p)

Return the roots of a polynomial with given coefficients. Wraps numpy.roots.

Source code in src/mechestim/_polynomial.py 
211
212
213
214
215
216
217
218
def roots(p):
    """Return the roots of a polynomial with given coefficients. Wraps ``numpy.roots``."""
    budget = require_budget()
    p = _np.asarray(p)
    n = len(p) - 1  # degree = number of roots
    cost = roots_cost(n)
    budget.deduct("roots", flop_cost=cost, subscripts=None, shapes=(p.shape,))
    return _np.roots(p)

roots_cost(n)

Cost for roots: \(10n^3\) FLOPs (companion matrix eigendecomposition).

Source code in src/mechestim/_polynomial.py 
60
61
62
def roots_cost(n: int) -> int:
    """Cost for roots: $10n^3$ FLOPs (companion matrix eigendecomposition)."""
    return max(10 * n**3, 1)

mechestim._window

Window function wrappers with FLOP counting.

bartlett_cost(n)

FLOP cost of Bartlett window generation.

Parameters:

Name Type Description Default
n int

Window length.

 required

Returns:

Type Description
int

Estimated FLOP count: n.
Notes

One linear evaluation per sample.

Source code in src/mechestim/_window.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def bartlett_cost(n: int) -> int:
    """FLOP cost of Bartlett window generation.

    Parameters
    ----------
    n : int
        Window length.

    Returns
    -------
    int
        Estimated FLOP count: n.

    Notes
    -----
    One linear evaluation per sample.
    """
    return max(n, 1)

blackman_cost(n)

FLOP cost of Blackman window generation.

Parameters:

Name Type Description Default
n int

Window length.

 required

Returns:

Type Description
int

Estimated FLOP count: 3n.
Notes

Three cosine terms per sample.

Source code in src/mechestim/_window.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def blackman_cost(n: int) -> int:
    """FLOP cost of Blackman window generation.

    Parameters
    ----------
    n : int
        Window length.

    Returns
    -------
    int
        Estimated FLOP count: 3n.

    Notes
    -----
    Three cosine terms per sample.
    """
    return max(3 * n, 1)

hamming_cost(n)

FLOP cost of Hamming window generation.

Parameters:

Name Type Description Default
n int

Window length.

 required

Returns:

Type Description
int

Estimated FLOP count: 2n.
Notes

One cosine evaluation plus one linear combination per sample.

Source code in src/mechestim/_window.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def hamming_cost(n: int) -> int:
    """FLOP cost of Hamming window generation.

    Parameters
    ----------
    n : int
        Window length.

    Returns
    -------
    int
        Estimated FLOP count: 2n.

    Notes
    -----
    One cosine evaluation plus one linear combination per sample.
    """
    return max(2 * n, 1)

hanning_cost(n)

FLOP cost of Hanning window generation.

Parameters:

Name Type Description Default
n int

Window length.

 required

Returns:

Type Description
int

Estimated FLOP count: 2n.
Notes

One cosine evaluation plus one linear combination per sample.

Source code in src/mechestim/_window.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def hanning_cost(n: int) -> int:
    """FLOP cost of Hanning window generation.

    Parameters
    ----------
    n : int
        Window length.

    Returns
    -------
    int
        Estimated FLOP count: 2n.

    Notes
    -----
    One cosine evaluation plus one linear combination per sample.
    """
    return max(2 * n, 1)

kaiser_cost(n)

FLOP cost of Kaiser window generation.

Parameters:

Name Type Description Default
n int

Window length.

 required

Returns:

Type Description
int

Estimated FLOP count: 3n.
Notes

Bessel function evaluation per sample.

Source code in src/mechestim/_window.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def kaiser_cost(n: int) -> int:
    """FLOP cost of Kaiser window generation.

    Parameters
    ----------
    n : int
        Window length.

    Returns
    -------
    int
        Estimated FLOP count: 3n.

    Notes
    -----
    Bessel function evaluation per sample.
    """
    return max(3 * n, 1)

mechestim._unwrap

Unwrap wrapper with FLOP counting.

unwrap_cost(shape)

FLOP cost of phase unwrapping.

Parameters:

Name Type Description Default
shape tuple of int

Input array shape.

 required

Returns:

Type Description
int

Estimated FLOP count: numel(input).
Notes

Cost covers element-wise differencing and conditional adjustment.

Source code in src/mechestim/_unwrap.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def unwrap_cost(shape: tuple[int, ...]) -> int:
    """FLOP cost of phase unwrapping.

    Parameters
    ----------
    shape : tuple of int
        Input array shape.

    Returns
    -------
    int
        Estimated FLOP count: numel(input).

    Notes
    -----
    Cost covers element-wise differencing and conditional adjustment.
    """
    numel = 1
    for d in shape:
        numel *= d
    return max(numel, 1)

mechestim._sorting_ops

Counted wrappers for sorting, search, and set operations.

argpartition(a, kth, axis=-1, **kwargs)

Counted version of numpy.argpartition. Cost: n FLOPs per slice.

Source code in src/mechestim/_sorting_ops.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def argpartition(a, kth, axis=-1, **kwargs):
    """Counted version of ``numpy.argpartition``. Cost: n FLOPs per slice."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if a.ndim == 0:
        cost = 1
    else:
        ax = axis if axis is not None else -1
        ax = ax % a.ndim
        n = a.shape[ax]
        numel = 1
        for d in a.shape:
            numel *= d
        num_slices = numel // n if n > 0 else 1
        cost = max(num_slices * n, 1)
    budget.deduct("argpartition", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    return _np.argpartition(a, kth, axis=axis, **kwargs)

argsort(a, axis=-1, **kwargs)

Counted version of numpy.argsort. Cost: n*ceil(log2(n)) FLOPs per slice.

Source code in src/mechestim/_sorting_ops.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def argsort(a, axis=-1, **kwargs):
    """Counted version of ``numpy.argsort``. Cost: n*ceil(log2(n)) FLOPs per slice."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if a.ndim == 0:
        cost = 1
    elif axis is None:
        cost = sort_cost(a.size)
    else:
        ax = axis % a.ndim
        cost = _sort_cost_nd(a, ax)
    budget.deduct("argsort", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    return _np.argsort(a, axis=axis, **kwargs)

digitize(x, bins, **kwargs)

Counted version of numpy.digitize.

Cost: n * ceil(log2(len(bins))) where n = numel(x).

Source code in src/mechestim/_sorting_ops.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def digitize(x, bins, **kwargs):
    """Counted version of ``numpy.digitize``.

    Cost: n * ceil(log2(len(bins))) where n = numel(x).
    """
    budget = require_budget()
    x_arr = _np.asarray(x)
    bins_arr = _np.asarray(bins)
    n = max(x_arr.size, 1)
    cost = search_cost(n, max(len(bins_arr), 1))
    budget.deduct(
        "digitize",
        flop_cost=cost,
        subscripts=None,
        shapes=(x_arr.shape, bins_arr.shape),
    )
    return _np.digitize(x, bins, **kwargs)

in1d(ar1, ar2, **kwargs)

Counted version of numpy.in1d. Cost: (n+m)*ceil(log2(n+m)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
307
308
309
310
311
312
313
314
def in1d(ar1, ar2, **kwargs):
    """Counted version of ``numpy.in1d``. Cost: (n+m)*ceil(log2(n+m)) FLOPs."""
    budget = require_budget()
    a1 = _np.asarray(ar1)
    a2 = _np.asarray(ar2)
    cost = _set_cost(a1, a2)
    budget.deduct("in1d", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape))
    return _np.in1d(ar1, ar2, **kwargs)

intersect1d(ar1, ar2, **kwargs)

Counted version of numpy.intersect1d. Cost: (n+m)*ceil(log2(n+m)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
333
334
335
336
337
338
339
340
341
342
def intersect1d(ar1, ar2, **kwargs):
    """Counted version of ``numpy.intersect1d``. Cost: (n+m)*ceil(log2(n+m)) FLOPs."""
    budget = require_budget()
    a1 = _np.asarray(ar1)
    a2 = _np.asarray(ar2)
    cost = _set_cost(a1, a2)
    budget.deduct(
        "intersect1d", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape)
    )
    return _np.intersect1d(ar1, ar2, **kwargs)

isin(element, test_elements, **kwargs)

Counted version of numpy.isin. Cost: (n+m)*ceil(log2(n+m)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
320
321
322
323
324
325
326
327
def isin(element, test_elements, **kwargs):
    """Counted version of ``numpy.isin``. Cost: (n+m)*ceil(log2(n+m)) FLOPs."""
    budget = require_budget()
    el = _np.asarray(element)
    te = _np.asarray(test_elements)
    cost = _set_cost(el, te)
    budget.deduct("isin", flop_cost=cost, subscripts=None, shapes=(el.shape, te.shape))
    return _np.isin(element, test_elements, **kwargs)

lexsort(keys, axis=-1)

Counted version of numpy.lexsort.

Cost: k * sort_cost(n) where k = len(keys) and n = len(keys[0]).

Source code in src/mechestim/_sorting_ops.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def lexsort(keys, axis=-1):
    """Counted version of ``numpy.lexsort``.

    Cost: k * sort_cost(n) where k = len(keys) and n = len(keys[0]).
    """
    budget = require_budget()
    # keys is a sequence of arrays; convert to list for inspection
    keys_list = list(keys)
    k = len(keys_list)
    if k == 0:
        cost = 1
    else:
        # numpy.lexsort uses the last key as primary; length is shape along axis
        first = _np.asarray(keys_list[0])
        n = first.shape[axis] if first.ndim > 0 else 1
        cost = max(k * sort_cost(n), 1)
    shapes = tuple(_np.asarray(key).shape for key in keys_list)
    budget.deduct("lexsort", flop_cost=cost, subscripts=None, shapes=shapes)
    return _np.lexsort(keys_list, axis=axis)

partition(a, kth, axis=-1, **kwargs)

Counted version of numpy.partition. Cost: n FLOPs per slice (quickselect).

Source code in src/mechestim/_sorting_ops.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def partition(a, kth, axis=-1, **kwargs):
    """Counted version of ``numpy.partition``. Cost: n FLOPs per slice (quickselect)."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if a.ndim == 0:
        cost = 1
    else:
        ax = axis if axis is not None else -1
        ax = ax % a.ndim
        n = a.shape[ax]
        numel = 1
        for d in a.shape:
            numel *= d
        num_slices = numel // n if n > 0 else 1
        cost = max(num_slices * n, 1)
    budget.deduct("partition", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    return _np.partition(a, kth, axis=axis, **kwargs)

searchsorted(a, v, **kwargs)

Counted version of numpy.searchsorted.

Cost: m * ceil(log2(n)) where m = numel(v) and n = len(a).

Source code in src/mechestim/_sorting_ops.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def searchsorted(a, v, **kwargs):
    """Counted version of ``numpy.searchsorted``.

    Cost: m * ceil(log2(n)) where m = numel(v) and n = len(a).
    """
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    v_arr = _np.asarray(v)
    n = a.shape[0] if a.ndim > 0 else 1
    m = max(v_arr.size, 1)
    cost = search_cost(m, n)
    budget.deduct(
        "searchsorted",
        flop_cost=cost,
        subscripts=None,
        shapes=(a.shape, v_arr.shape),
    )
    return _np.searchsorted(a, v, **kwargs)

setdiff1d(ar1, ar2, **kwargs)

Counted version of numpy.setdiff1d. Cost: (n+m)*ceil(log2(n+m)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
365
366
367
368
369
370
371
372
373
374
def setdiff1d(ar1, ar2, **kwargs):
    """Counted version of ``numpy.setdiff1d``. Cost: (n+m)*ceil(log2(n+m)) FLOPs."""
    budget = require_budget()
    a1 = _np.asarray(ar1)
    a2 = _np.asarray(ar2)
    cost = _set_cost(a1, a2)
    budget.deduct(
        "setdiff1d", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape)
    )
    return _np.setdiff1d(ar1, ar2, **kwargs)

setxor1d(ar1, ar2, **kwargs)

Counted version of numpy.setxor1d. Cost: (n+m)*ceil(log2(n+m)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
382
383
384
385
386
387
388
389
390
391
def setxor1d(ar1, ar2, **kwargs):
    """Counted version of ``numpy.setxor1d``. Cost: (n+m)*ceil(log2(n+m)) FLOPs."""
    budget = require_budget()
    a1 = _np.asarray(ar1)
    a2 = _np.asarray(ar2)
    cost = _set_cost(a1, a2)
    budget.deduct(
        "setxor1d", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape)
    )
    return _np.setxor1d(ar1, ar2, **kwargs)

sort(a, axis=-1, **kwargs)

Counted version of numpy.sort. Cost: n*ceil(log2(n)) FLOPs per slice.

Source code in src/mechestim/_sorting_ops.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def sort(a, axis=-1, **kwargs):
    """Counted version of ``numpy.sort``. Cost: n*ceil(log2(n)) FLOPs per slice."""
    budget = require_budget()
    if not isinstance(a, _np.ndarray):
        a = _np.asarray(a)
    if a.ndim == 0:
        cost = 1
    elif axis is None:
        cost = sort_cost(a.size)
    else:
        ax = axis % a.ndim
        cost = _sort_cost_nd(a, ax)
    budget.deduct("sort", flop_cost=cost, subscripts=None, shapes=(a.shape,))
    return _np.sort(a, axis=axis, **kwargs)

union1d(ar1, ar2)

Counted version of numpy.union1d. Cost: (n+m)*ceil(log2(n+m)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
350
351
352
353
354
355
356
357
358
359
def union1d(ar1, ar2):
    """Counted version of ``numpy.union1d``. Cost: (n+m)*ceil(log2(n+m)) FLOPs."""
    budget = require_budget()
    a1 = _np.asarray(ar1)
    a2 = _np.asarray(ar2)
    cost = _set_cost(a1, a2)
    budget.deduct(
        "union1d", flop_cost=cost, subscripts=None, shapes=(a1.shape, a2.shape)
    )
    return _np.union1d(ar1, ar2)

unique(ar, **kwargs)

Counted version of numpy.unique. Cost: n*ceil(log2(n)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
222
223
224
225
226
227
228
def unique(ar, **kwargs):
    """Counted version of ``numpy.unique``. Cost: n*ceil(log2(n)) FLOPs."""
    budget = require_budget()
    ar_arr = _np.asarray(ar)
    cost = _unique_cost(ar_arr)
    budget.deduct("unique", flop_cost=cost, subscripts=None, shapes=(ar_arr.shape,))
    return _np.unique(ar_arr, **kwargs)

unique_all(x)

Counted version of numpy.unique_all. Cost: n*ceil(log2(n)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
234
235
236
237
238
239
240
def unique_all(x, /):
    """Counted version of ``numpy.unique_all``. Cost: n*ceil(log2(n)) FLOPs."""
    budget = require_budget()
    x_arr = _np.asarray(x)
    cost = _unique_cost(x_arr)
    budget.deduct("unique_all", flop_cost=cost, subscripts=None, shapes=(x_arr.shape,))
    return _np.unique_all(x_arr)

unique_counts(x)

Counted version of numpy.unique_counts. Cost: n*ceil(log2(n)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
246
247
248
249
250
251
252
253
254
def unique_counts(x, /):
    """Counted version of ``numpy.unique_counts``. Cost: n*ceil(log2(n)) FLOPs."""
    budget = require_budget()
    x_arr = _np.asarray(x)
    cost = _unique_cost(x_arr)
    budget.deduct(
        "unique_counts", flop_cost=cost, subscripts=None, shapes=(x_arr.shape,)
    )
    return _np.unique_counts(x_arr)

unique_inverse(x)

Counted version of numpy.unique_inverse. Cost: n*ceil(log2(n)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
262
263
264
265
266
267
268
269
270
def unique_inverse(x, /):
    """Counted version of ``numpy.unique_inverse``. Cost: n*ceil(log2(n)) FLOPs."""
    budget = require_budget()
    x_arr = _np.asarray(x)
    cost = _unique_cost(x_arr)
    budget.deduct(
        "unique_inverse", flop_cost=cost, subscripts=None, shapes=(x_arr.shape,)
    )
    return _np.unique_inverse(x_arr)

unique_values(x)

Counted version of numpy.unique_values. Cost: n*ceil(log2(n)) FLOPs.

Source code in src/mechestim/_sorting_ops.py 
278
279
280
281
282
283
284
285
286
def unique_values(x, /):
    """Counted version of ``numpy.unique_values``. Cost: n*ceil(log2(n)) FLOPs."""
    budget = require_budget()
    x_arr = _np.asarray(x)
    cost = _unique_cost(x_arr)
    budget.deduct(
        "unique_values", flop_cost=cost, subscripts=None, shapes=(x_arr.shape,)
    )
    return _np.unique_values(x_arr)

mechestim._counting_ops

Counted wrappers for trace, histogram, and generation operations.

These are operations that look "free" but involve genuine computation. Each function charges a FLOP cost to the active budget.