From e3bbe3cb9e173e6547e54e67a9ab7d7f8dfd5fee Mon Sep 17 00:00:00 2001 From: Jeffrey Ward Date: Sat, 11 Jul 2026 14:57:34 -0400 Subject: [PATCH 1/2] feat: scale a batched tensor by a vector of per-batch scalars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- src/tensors/base_tensor/base_tensor.jl | 15 ++++++++++++ test/test_pyclasses.jl | 33 +++++++++++++++++++++++--- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/src/tensors/base_tensor/base_tensor.jl b/src/tensors/base_tensor/base_tensor.jl index 91884d0..fadb8fc 100644 --- a/src/tensors/base_tensor/base_tensor.jl +++ b/src/tensors/base_tensor/base_tensor.jl @@ -24,6 +24,21 @@ Base.:*(s::Number, t::AbstractTensor) = wrap(t.space, s .* t.coeffs) Base.:*(t::AbstractTensor, s::Number) = s * t Base.:/(t::AbstractTensor, s::Number) = wrap(t.space, t.coeffs ./ s) +# ── Per-batch scalar weights ─────────────────────────────────────────────────── +# A plain numeric array is a bag of *batch-wise* scalars, never a coefficient +# vector: its shape broadcasts (numpy right-alignment) against the tensor's batch +# axes and the coefficient axis is left untouched. This is what makes a spectral +# filter `exp.(-vals) * vecs` work on the batched eigenbasis from `spectrum`. +function _scale_by_batch(t::AbstractTensor, w::AbstractArray{<:Number}, op) + @assert compatible_batches(batch_shape(t), size(w)) "Incompatible batch shapes: $(batch_shape(t)) vs $(size(w))" + ct, wt = align_batch_pair(t.coeffs, reshape(w, size(w)..., 1)) + return wrap(t.space, op.(ct, wt)) +end + +Base.:*(w::AbstractArray{<:Number}, t::AbstractTensor) = _scale_by_batch(t, w, *) +Base.:*(t::AbstractTensor, w::AbstractArray{<:Number}) = _scale_by_batch(t, w, *) +Base.:/(t::AbstractTensor, w::AbstractArray{<:Number}) = _scale_by_batch(t, w, /) + # A scalar added to a function is the constant function of that value, projected # onto the basis. Only functions absorb a scalar this way — the other tensors have # no canonical constant element, and `+`/`-` against a Number stays a MethodError. diff --git a/test/test_pyclasses.jl b/test/test_pyclasses.jl index 7c4b071..7e00444 100644 --- a/test/test_pyclasses.jl +++ b/test/test_pyclasses.jl @@ -774,6 +774,19 @@ end @test vals_only isa AbstractVector @test all(≈(1.0), vals_only) end + + # A spectral filter: weight each eigenform by a function of its eigenvalue. + @testset "eigenvector batch scaling by a weight vector" begin + vals, vecs = spectrum(up_laplacian(dg, 1) + 1.0 * down_laplacian(dg, 1)) + weights = exp.(-vals) + expected = coeffs(vecs) .* reshape(weights, batch_shape(vecs)..., 1) + + for out in (weights * vecs, vecs * weights) + @test typeof(out) === typeof(vecs) + @test batch_shape(out) == batch_shape(vecs) + @test coeffs(out) ≈ expected + end + end end # Skipped: `repr(op)` contents (Julia's `show` for LinearOperator has a different @@ -878,14 +891,28 @@ end @test coeffs(ωb * 3.0) ≈ cb .* 3.0 end - # GAP: Python scales a batched tensor by a length-B vector of per-batch - # scalars (`weights * omega`). Julia's `*` only accepts a Number. @testset "per-batch scalar vector scaling" begin B = 5 cb = randn(Xoshiro(0), B, n1 * dim) ωb = wrap(f1s, cb) weights = exp.(-range(0.2, 1.2; length=B)) - @test_broken coeffs(weights * ωb) ≈ cb .* weights + + @test (weights * ωb) isa Form + @test coeffs(weights * ωb) ≈ cb .* weights + @test coeffs(ωb * weights) ≈ cb .* weights + @test coeffs(ωb / weights) ≈ cb ./ weights + + # The weights broadcast against the *batch* axes only, numpy-style + # (right-aligned), and can expand the batch shape. + c = randn(Xoshiro(1), n1 * dim) + @test coeffs(weights * wrap(f1s, c)) ≈ weights .* transpose(c) + + cb2 = randn(Xoshiro(2), 3, B, n1 * dim) + @test coeffs(weights * wrap(f1s, cb2)) ≈ cb2 .* reshape(weights, 1, B, 1) + + # Weights that no batch axis can absorb are an error, never a silent + # match against the coefficient axis. + @test_throws AssertionError randn(B + 1) * ωb end @testset "function-form product" begin From 220eb2e89d7b6190d81b3c0b2b9f09f4111a917e Mon Sep 17 00:00:00 2001 From: Jeffrey Ward Date: Sat, 11 Jul 2026 15:06:16 -0400 Subject: [PATCH 2/2] test: drop the stale @test_broken note, cover the error path for / --- test/test_pyclasses.jl | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/test/test_pyclasses.jl b/test/test_pyclasses.jl index 7e00444..c782b9e 100644 --- a/test/test_pyclasses.jl +++ b/test/test_pyclasses.jl @@ -4,11 +4,6 @@ # wrappers, batching/broadcasting, direct sums, spectral operators, transposes and # error paths. It runs over the same d = 1..4 config sweep (see pysuite.jl). # -# One upstream behaviour exercised here has no Julia implementation, and is -# recorded as @test_broken rather than dropped — it will flip to a failure the day -# someone adds it: scaling a batched tensor by a *vector* of per-batch scalars -# (Python's `weights * omega`, weights of length B). -# # Purely Python-shaped tests are skipped with a note where they appear: numpy ufunc # dispatch (`np.multiply(..., where=...)`), `repr` string contents (Julia defines # `show` only for LinearOperator), the numpy RuntimeWarning check, and the @@ -913,6 +908,7 @@ end # Weights that no batch axis can absorb are an error, never a silent # match against the coefficient axis. @test_throws AssertionError randn(B + 1) * ωb + @test_throws AssertionError ωb / randn(B + 1) end @testset "function-form product" begin