pymc-devs / pymc-devs/pytensor

DOC: document f.trust_input=True as MLX hot-loop performance tip

Open Beginner friendly
#2,093 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

docs mlx
Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

Part of #2085.

Current state

pytensor.compile.executor.Function.trust_input already exists and, when set to True, skips _validate_inputs entirely:

        if self.trust_input:
            for storage_data, arg in zip(self._input_storage_data, args):
                storage_data[0] = arg
            ...
        else:
            self._validate_inputs(args, kwargs)

For users who keep their data on device as mx.array and call the same compiled function in a hot loop (training, MCMC, batched inference), bypassing the type filter is a huge win — independent of #2087.

Metrics

f_mlx(mx_x, mx_y, mx_z) for sin(x)*exp(y)+z on 2048² (median, with mx.synchronize between calls):

Configuration Median (us) vs raw MLX
Raw mx.compile (lower bound) 478 1.0×
f_mlx, default 5854 12.2×
f_mlx, with #2087 filter fix 1397 2.9×
f_mlx, trust_input=True (no filter at all) 459 0.96×
f_mlx, trust_input=True + #2087 filter fix 459 0.96×

For workloads where the user has already validated input shapes/dtypes (e.g. typed function call inside a sampler step), trust_input=True puts PyTensor MLX at parity with hand-written mx.compile.

Other workload examples (mx input, post all proposed fixes):

Workload trust_input=False trust_input=True
Variadic add 5×512² 310 µs 237 µs
Argmax axis=1 on (128, 1024, 16) 555 µs 330 µs
MLP forward+grad 1176 µs 997 µs

Proposed change

Documentation only. Add a small "Performance" section to the MLX backend docs (or to doc/links.rst or wherever the MLX page lives) noting:

  1. PyTensor performs an _validate_inputs step on every call that converts non-np.ndarray array inputs through np.asarray(...) round-trips. The fast path being added in #2087 covers most cases.
  2. For hot-loop inference / sampling where the user controls the input types, f.trust_input = True skips this validation entirely and brings PyTensor MLX to within ~3 % of raw mx.compile performance.
  3. Caveat: with trust_input=True, the user must pass values whose types and shapes already match the compiled function's expectations. Wrong dtype / shape will not be caught and may produce wrong results or MLX errors deeper in the pipeline.
Suggested snippet for the docs
.. _mlx_perf_tips:

Performance tips for the MLX backend
====================================

Skip input validation in hot loops
----------------------------------

PyTensor validates input types and shapes on every call to a compiled
function. This adds dispatch overhead that's negligible for slow ops but
visible for short MLX kernels. If you're calling the same function many
times with already-correct ``mx.array`` inputs (e.g. inside a training step
or MCMC iteration), set ``f.trust_input = True`` after compilation:

.. code-block:: python

    f = pytensor.function([x, y], expr, mode="MLX")
    f.trust_input = True   # skip input validation
    for _ in range(n_steps):
        out = f(x_mx, y_mx)
        ...

Benchmarks at 2048×2048 elementwise show this brings PyTensor MLX overhead
from roughly 3× raw ``mx.compile`` down to under 5 %. The trade-off is no
runtime check that arguments match the compiled signature; pass wrong dtypes
or shapes at your own risk.

Acceptance criteria

  • doc/ page updated with the section above (or equivalent).
  • Cross-reference from the MLX backend overview, if one exists.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by locating the MLX backend overview under doc/; the issue also mentions doc/links.rst as a possible location. Add a Performance section covering trust_input=True, its hot-loop benefit, the required input-type and shape caveat, and the provided usage example. Done means the relevant doc page is updated and cross-referenced from the MLX overview when one exists.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
documentation, performance
Issue type
Documentation
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
76/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.