Skip to content

feat: scale a batched tensor by a vector of per-batch scalars - #20

Merged
wardjm merged 2 commits into
mainfrom
batch-scalar-vector-scaling
Jul 11, 2026
Merged

wardjm merged 2 commits into
mainfrom
batch-scalar-vector-scaling

Conversation

@wardjm

@wardjm wardjm commented Jul 11, 2026 •

Copy link
Copy Markdown
Owner

Closes #12.

A plain numeric array multiplying or dividing a tensor is now read as batch-wise scalar weights: weights * ω, ω * weights and ω / weights all work. The weights' shape broadcasts against the tensor's batch axes with numpy's right-alignment (reusing compatible_batches / align_batch_pair), and the coefficient axis is never matched against them. Weights may also expand the batch shape, so an unbatched form times a length-B vector yields batch (B,) — the same rule as upstream's _broadcast_batch_scalars + _broadcast_coeffs_to_batch.

This makes a spectral filter work as it does in Python:

vals, vecs = spectrum(up_laplacian(dg, 1) + down_laplacian(dg, 1))
exp.(-vals) * vecs      # weight each eigenform by a function of its eigenvalue

Tests

The @test_broken in test/test_pyclasses.jl (product_operations → per-batch scalar vector scaling) is flipped to real assertions and extended to cover the reverse order, division, batch expansion, multi-dimensional batch broadcasting, and the error path when no batch axis can absorb the weights. The upstream test that motivated the issue, test_operators.py::test_spectrum_eigenvector_batch_scaling_with_numpy, is ported into the operators section.

Full suite passes: 2381 tests, none broken.

Two deliberate divergences

The dispatch signature is AbstractArray{<:Number}, not the AbstractVector the issue suggested. Multi-dimensional batch shapes need the wider signature and upstream's np.broadcast_shapes accepts any rank — but it does mean a matrix times a tensor is now a batch-broadcast scale rather than a MethodError, so Julia's usual linear-algebra reading of A * x doesn't apply here.

Incompatible weights raise AssertionError, where upstream raises TypeError. That matches how + and the pointwise products already guard in this file, so it is consistent with the repo rather than with Python.

Not covered

weights .* ω (broadcast syntax) still fails — AbstractTensor has no Base.broadcastable. Upstream reaches it via __array_ufunc__, which this test file already declares out of scope. Filed separately as #19.

wardjm added 2 commits July 11, 2026 14:57
A plain numeric array multiplying (or dividing) a tensor is now read as
batch-wise scalar weights: its shape broadcasts against the batch axes with
numpy's right-alignment, and the coefficient axis is never matched against it.
This makes a spectral filter — exp.(-vals) * vecs on the batched eigenbasis
from spectrum — work as it does upstream.

Closes #12
@wardjm
wardjm merged commit 42617a0 into main Jul 11, 2026
2 checks passed
@wardjm
wardjm deleted the batch-scalar-vector-scaling branch August 3, 2026 20:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Cannot scale a batched tensor by a vector of per-batch scalars

1 participant