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

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.

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
@@ -774,6 +769,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 +886,29 @@ 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
+            @test_throws AssertionError ωb / randn(B + 1)
         end
 
         @testset "function-form product" begin