[dexhunter]

TL;DR

• Speeds up fast_grad in pina/_src/core/operator.py, which pina.operator.grad calls for every gradient.
• Slices the plain output tensor, turns off LabelTensor dispatch around autograd.grad, and skips unneeded index and cat copies.
• Benchmark (workload like tests/test_operator.py::test_gradient, median of 3 timed runs on an isolated core): 147.9 ms -> 98.9 ms, 33% less time.
• Tested with the operator, equation, solver, problem zoo and label tensor tests: 664 passed.

Refs #854

Description

This PR fixes #854. It changes only fast_grad and its helper _scalar_grad.

What & Why

Before, fast_grad called LabelTensor.extract for each output component. Every torch call on a LabelTensor also went through __torch_function__. _scalar_grad looked up the input label indices again for each component. It also copied the gradient with an index, even when d asked for all inputs in order.

Now fast_grad takes the plain tensor once and slices each component by its column index. It computes the input label indices once and passes them to _scalar_grad. It skips the index when d equals the input labels, and skips torch.cat for one component. _scalar_grad runs torch.autograd.grad with the dispatch turned off, so the gradient comes back as a plain tensor. The math and the input object stay the same, so values, labels and higher-order gradients do not change.

Changes

• fast_grad: use one shared path for scalar and vector outputs, with column slices and precomputed label indices.
• _scalar_grad: take column indices (or None for all inputs) and return a plain tensor.

Results

Python 3.12.13, torch 2.15.0+cpu, Linux x86_64, one thread, isolated core, best of 7 runs, median of 3 rounds:

Situation	Before	After	Change
grad, 1 input, 2 outputs, 20 rows	365 µs	220 µs	1.66x faster
grad, 1 input, 1 output, 20 rows	125 µs	94 µs	1.33x faster
grad, 2 inputs, 2 outputs, 20 rows	392 µs	233 µs	1.68x faster
grad, 3 inputs, 3 outputs, d=["x", "z"], 1000 rows	719 µs	520 µs	1.38x faster
grad, 1 input, 2 outputs, 1000000 rows (math dominates, small gain)	45.2 ms	39.8 ms	12% faster, no slowdown
Workload like test_gradient (mostly 20-row, 2-output cases)	147.9 ms	98.9 ms	33% less time

Why this is safe and easy to maintain

The diff is small and replaces the two scalar/vector branches with one shorter path. Gradient values, dtypes, shapes, labels and create_graph=True behaviour stay the same. tests/test_operator.py checks grad, div, laplacian and advection, also for higher-order derivatives, and all tests pass.

Validation

• Benchmark tests: pytest tests/test_label_tensor/test_label_tensor.py tests/test_condition/test_input_target_condition.py tests/test_operator.py tests/test_condition/test_data_condition.py -> 103 passed.
• Wider tests: pytest tests/test_operator.py tests/test_equation tests/test_equation_zoo tests/test_solver tests/test_problem_zoo tests/test_label_tensor -> 664 passed.
• Code style: black --line-length 80 leaves the file unchanged.
• Short reproduction: python -m timeit -s "import torch; from pina import LabelTensor; from pina.operator import grad; x = LabelTensor(torch.rand(20, 1), ['x']).requires_grad_(); u = LabelTensor(torch.cat([x**2, x**3], -1), ['u', 'v'])" "grad(u, x)"

Checklist

• Code follows the project’s Code Style Guidelines
• Tests have been added or updated
• Documentation has been updated if necessary
• Pull request is linked to an open issue

AI assistance

OpenAI Codex (gpt-6.1-sol) wrote the candidate versions of these two functions. Weco Observe tracked and scored them. A Claude Code agent picked a short, readable version that was clearly faster, and ran the tests and benchmarks above.

Weco share link: https://dashboard.weco.ai/share/5TM-D6CrkwpPBZFoFVh6IvdS7hwY-6M6

[GiovanniCanali]

Thank you, @dexhunter! This looks great. Excellent catch identifying the performance bottleneck caused by the unnecessary use of LabelTensor in the _scalar_grad utility method.

I believe this improvement will have a significant impact, particularly on physics-informed models, where gradients are computed extensively during each training epoch.

[GiovanniCanali]

Just a quick note, @dexhunter: I've changed the target branch for this PR from master to dev, as we typically merge contributions into dev first and periodically release them to master in batches, alongside other changes.

For future development and contributions to the package, we recommend using the dev branch as your starting point.