Skip to content

Accept any AbstractVecOrMat when sampling from an MvNormal - #2087

Merged
devmotion merged 2 commits into
masterfrom
mvnormal-rand-abstractvecormat
Aug 20, 2026
Merged

devmotion merged 2 commits into
masterfrom
mvnormal-rand-abstractvecormat

Conversation

@devmotion

@devmotion devmotion commented Aug 19, 2026 •

Copy link
Copy Markdown
Member

Fixes #2086.

_rand! for MvNormal was defined for VecOrMat === Union{Vector,Matrix}, so any other AbstractMatrix — a view, most commonly — missed it and fell back to the generic eachvariate method in genericrand.jl, which draws column by column. That replaced one randn! and one BLAS-3 unwhiten! with size(x, 2) scalar randn loops and size(x, 2) BLAS-2 unwhiten! calls.

This widens the signature to AbstractVecOrMat and drops the separate scalar-loop method for AbstractVector, whose comment — "randn! only works for Array, but not generally for AbstractArray" — no longer holds: randn! has a generic AbstractArray method on every Julia version this package supports (julia = "1.10"), and MvNormalCanon and GenericMvTDist already accept an AbstractMatrix and rely on it.

For a 10×1000 draw with a dense covariance matrix:

before after
MvNormal, Matrix 21.416 μs 21.416 μs
MvNormal, view 115.583 μs 43.375 μs
MvLogNormal, view 149.625 μs 76.417 μs

MvLogNormal follows because it takes an AbstractVecOrMat and forwards into MvNormal. The view path now sits level with MvNormalCanon's (45.875 μs), which already had the wider signature.

Behaviour:

  • Vector and Matrix arguments keep the body they already dispatched to; a seeded Matrix draw hashes identically before and after.
  • For other AbstractVecOrMat the drawn normals are unchanged, since the generic randn! walks eachindex in the same column-major order the per-column loop did. Only unwhiten! moves from size(x, 2) BLAS-2 calls to one BLAS-3 call, which is bit-identical for a PDiagMat and differs at the rounding level for a PDMat (maxdiff = 8.9e-16 on a 6×40 draw).
  • Test.detect_ambiguities(Distributions; recursive = true) stays at 0. (MvNormal, AbstractVecOrMat) is strictly more specific than (Sampleable{<:ArrayLikeVariate}, AbstractArray) in both arguments.
  • PDMats' unwhiten! is already declared for AbstractVecOrMat throughout, so no companion change is needed.
  • MvLogitNormal is affected beyond the view case. Its _rand! passes @views _drop1(x) into rand!(rng, d.normal, …), i.e. always a SubArray, so even a plain Matrix draw moves from the column loop to the bulk path and its output changes at the rounding level in the common case.

Since the change is behaviour-preserving up to rounding, the testset can only do so much. The one assertion that fails on master is that sampling into a view of a whole matrix gives exactly the same samples as sampling into the matrix, and it fails there through the rounding difference between one BLAS-3 unwhiten! and one BLAS-2 call per column (maxdiff = 4.4e-16 for the covariance matrix used). The remaining assertions guard against breakage rather than proving that the bulk path is taken: a view of some columns is filled like a matrix of the same size and nothing outside it is written, for a PDMat, a PDiagMat and a ScalMat, for a view that is not strided and hence cannot use BLAS in unwhiten!, for a single column in place of the removed AbstractVector method, and for a view without columns.

That exact equality holds only because StableRNG has no bulk randn!. The ones in Random are defined for MersenneTwister/Xoshiro and Array{Float64}, and a view satisfies neither, so with those RNGs a Matrix and a view draw different streams by design, both before and after this PR. This is also the residual performance gap between the two, and it is a Julia-level matter that is not addressed here.


This pull request was written with the assistance of generative AI (Claude Code).

`_rand!` was defined for `VecOrMat === Union{Vector,Matrix}`, so any other
`AbstractMatrix` such as a `view` fell back to the generic `eachvariate`
method in `genericrand.jl` and was drawn column by column. That replaced one
`randn!` and one BLAS-3 `unwhiten!` with `size(x, 2)` scalar `randn` loops and
`size(x, 2)` BLAS-2 `unwhiten!` calls, which for a 10x1000 view cost 115.6us
against 21.4us for the equivalent `Matrix`.

`randn!` has a generic `AbstractArray` method on every supported Julia version,
so the separate scalar-loop method for `AbstractVector` and the comment
justifying it are no longer needed. `MvNormalCanon` and `GenericMvTDist`
already accept an `AbstractMatrix` and rely on that generic `randn!`.

`Vector` and `Matrix` arguments keep the body they already dispatched to. For
other `AbstractVecOrMat` the drawn normals are unchanged, since the generic
`randn!` walks `eachindex` in the same column-major order the per-column loop
did; only `unwhiten!` moves from `size(x, 2)` BLAS-2 calls to one BLAS-3 call,
which is bit-identical for a `PDiagMat` and differs at the rounding level for
a `PDMat`.

Fixes #2086.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@codecov

codecov Bot commented Aug 19, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.76%. Comparing base (54720dc) to head (0d6ea6b).

Additional details and impacted files
@@            Coverage Diff             @@
##           master    #2087      +/-   ##
==========================================
+ Coverage   86.75%   86.76%   +0.01%     
==========================================
  Files         149      149              
  Lines        8881     8888       +7     
==========================================
+ Hits         7705     7712       +7     
  Misses       1176     1176              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@devmotion
devmotion requested a review from andreasnoack August 20, 2026 07:18
Three of the four previous assertions passed unchanged on master. The
per-column comparison asserts the `eachindex` ordering of `Random.randn!`
rather than anything about this change, and the comparison of a column-range
view against a slice of a larger draw holds either way. Only the exact
comparison of a view of a whole matrix against the matrix fails without this
change, and it does so through the rounding difference between one BLAS-3
`unwhiten!` and one BLAS-2 call per column, 4.4e-16 for the covariance matrix
used here.

The comment claiming that a view takes the same code path as a matrix was
therefore misleading. The bulk `randn!` methods in `Random` are defined for
`MersenneTwister`/`Xoshiro` and `Array{Float64}`, and a view satisfies
neither, so the samples are exactly equal only for an RNG such as `StableRNG`
that has no bulk method. The comment now says so, since a future `randn!` for
any `AbstractRNG` would break the assertion with this change applied.

The testset also covered a single dense `PDMat` and, apart from the per-column
loop, only contiguous views. It now runs for a `PDMat`, a `PDiagMat` and a
`ScalMat`, whose `unwhiten!` is a broadcast and hence exact on any BLAS. It
additionally covers a view that is not strided, for which `unwhiten!` cannot
use BLAS, a view of a single column in place of the removed `AbstractVector`
method, and a view without columns, where the old fallback loop did nothing
and a single `unwhiten!` without columns is made instead.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@devmotion
devmotion merged commit 050e0f9 into master Aug 20, 2026
14 checks passed
@devmotion
devmotion deleted the mvnormal-rand-abstractvecormat branch August 20, 2026 11:42
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.

rand! for MvNormal is much slower for an AbstractMatrix than for a Matrix

2 participants