pymc-devs / pymc-devs/pytensor
DOC: document f.trust_input=True as MLX hot-loop performance tip
Nobody has claimed this yet.
- 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:
- PyTensor performs an
_validate_inputsstep on every call that converts non-np.ndarrayarray inputs throughnp.asarray(...)round-trips. The fast path being added in #2087 covers most cases. - For hot-loop inference / sampling where the user controls the input types,
f.trust_input = Trueskips this validation entirely and brings PyTensor MLX to within ~3 % of rawmx.compileperformance. - 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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