add reinterpret(reshape, T, a) (#722)

Co-authored-by: Michael Abbott <me@escbook>

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "Compat"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "3.17.0"
+version = "3.18.0"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

README.md

@@ -55,6 +55,10 @@ changes in `julia`.
 
 ## Supported features
 
+* `reinterpret(reshape, T, a::AbstractArray{S})` reinterprets `a` to have eltype `T` while
+  inserting or consuming the first dimension depending on the ratio of `sizeof(T)` and `sizeof(S)`.
+  ([#37559]). (since Compat 3.18)
+
 * The composition operator `∘` now returns a `Compat.ComposedFunction` instead of an anonymous function ([#37517]). (since Compat 3.17)
 
 * New function `addenv` for adding environment mappings into a `Cmd` object, returning the new `Cmd` object ([#37244]). (since Compat 3.16)
@@ -203,3 +207,4 @@ Note that you should specify the correct minimum version for `Compat` in the
 [#35052]: https://github.com/JuliaLang/julia/pull/35052
 [#37244]: https://github.com/JuliaLang/julia/pull/37244
 [#37517]: https://github.com/JuliaLang/julia/pull/37517
+[#37559]: https://github.com/JuliaLang/julia/pull/37559

src/Compat.jl

@@ -663,6 +663,82 @@ if VERSION < v"1.6.0-DEV.873" # 18198b1bf85125de6cec266eac404d31ccc2e65c
     end
 end
 
+
+# https://github.com/JuliaLang/julia/pull/37559
+if VERSION < v"1.6.0-DEV.1083"
+    """
+        reinterpret(reshape, T, A::AbstractArray{S}) -> B
+
+    Change the type-interpretation of `A` while consuming or adding a "channel dimension."
+
+    If `sizeof(T) = n*sizeof(S)` for `n>1`, `A`'s first dimension must be
+    of size `n` and `B` lacks `A`'s first dimension. Conversely, if `sizeof(S) = n*sizeof(T)` for `n>1`,
+    `B` gets a new first dimension of size `n`. The dimensionality is unchanged if `sizeof(T) == sizeof(S)`.
+
+    # Examples
+
+    ```
+    julia> A = [1 2; 3 4]
+    2×2 Matrix{$Int}:
+     1  2
+     3  4
+
+    julia> reinterpret(reshape, Complex{Int}, A)    # the result is a vector
+    2-element reinterpret(reshape, Complex{$Int}, ::Matrix{$Int}):
+     1 + 3im
+     2 + 4im
+
+    julia> a = [(1,2,3), (4,5,6)]
+    2-element Vector{Tuple{$Int, $Int, $Int}}:
+     (1, 2, 3)
+     (4, 5, 6)
+
+    julia> reinterpret(reshape, Int, a)             # the result is a matrix
+    3×2 reinterpret(reshape, $Int, ::Vector{Tuple{$Int, $Int, $Int}}):
+     1  4
+     2  5
+     3  6
+    ```
+    """
+    function Base.reinterpret(::typeof(reshape), ::Type{T}, a::A) where {T,S,A<:AbstractArray{S}}
+        isbitstype(T) || throwbits(S, T, T)
+        isbitstype(S) || throwbits(S, T, S)
+        if sizeof(S) == sizeof(T)
+            N = ndims(a)
+        elseif sizeof(S) > sizeof(T)
+            rem(sizeof(S), sizeof(T)) == 0 || throwintmult(S, T)
+            N = ndims(a) + 1
+        else
+            rem(sizeof(T), sizeof(S)) == 0 || throwintmult(S, T)
+            N = ndims(a) - 1
+            N > -1 || throwsize0(S, T, "larger")
+            axes(a, 1) == Base.OneTo(sizeof(T) ÷ sizeof(S)) || throwsize1(a, T)
+        end
+        paxs = axes(a)
+        new_axes = if sizeof(S) > sizeof(T)
+            (Base.OneTo(div(sizeof(S), sizeof(T))), paxs...)
+        elseif sizeof(S) < sizeof(T)
+            Base.tail(paxs)
+        else
+            paxs
+        end
+        reshape(reinterpret(T, vec(a)), new_axes)
+    end
+
+    @noinline function throwintmult(S::Type, T::Type)
+        throw(ArgumentError("`reinterpret(reshape, T, a)` requires that one of `sizeof(T)` (got $(sizeof(T))) and `sizeof(eltype(a))` (got $(sizeof(S))) be an integer multiple of the other"))
+    end
+    @noinline function throwsize1(a::AbstractArray, T::Type)
+        throw(ArgumentError("`reinterpret(reshape, $T, a)` where `eltype(a)` is $(eltype(a)) requires that `axes(a, 1)` (got $(axes(a, 1))) be equal to 1:$(sizeof(T) ÷ sizeof(eltype(a))) (from the ratio of element sizes)"))
+    end
+    @noinline function throwbits(S::Type, T::Type, U::Type)
+        throw(ArgumentError("cannot reinterpret `$(S)` as `$(T)`, type `$(U)` is not a bits type"))
+    end
+    @noinline function throwsize0(S::Type, T::Type, msg)
+        throw(ArgumentError("cannot reinterpret a zero-dimensional `$(S)` array to `$(T)` which is of a $msg size"))
+    end
+end
+
 include("iterators.jl")
 include("deprecated.jl")
 

test/runtests.jl

@@ -600,6 +600,85 @@ end
     @test strip(String(read(cmd))) == "baz bar"
 end
 
+# https://github.com/JuliaLang/julia/pull/37559
+@testset "reinterpred(reshape, ...)" begin
+    # simplified from PR
+    Ar = Int64[1 3; 2 4]
+    @test @inferred(ndims(reinterpret(reshape, Complex{Int64}, Ar))) == 1
+    @test @inferred(axes(reinterpret(reshape, Complex{Int64}, Ar))) === (Base.OneTo(2),)
+    @test @inferred(size(reinterpret(reshape, Complex{Int64}, Ar))) == (2,)
+
+    _B = Complex{Int64}[5+6im, 7+8im, 9+10im]
+    @test @inferred(ndims(reinterpret(reshape, Int64, _B))) == 2
+    @test @inferred(axes(reinterpret(reshape, Int64, _B))) === (Base.OneTo(2), Base.OneTo(3))
+    @test @inferred(size(reinterpret(reshape, Int64, _B))) == (2, 3)
+    @test @inferred(ndims(reinterpret(reshape, Int128, _B))) == 1
+    @test @inferred(axes(reinterpret(reshape, Int128, _B))) === (Base.OneTo(3),)
+    @test @inferred(size(reinterpret(reshape, Int128, _B))) == (3,)
+
+    A = Int64[1, 2, 3, 4]
+    Av = [Int32[1,2], Int32[3,4]]
+
+    @test_throws ArgumentError reinterpret(Vector{Int64}, A) # ("cannot reinterpret `Int64` as `Vector{Int64}`, type `Vector{Int64}` is not a bits type")
+    @test_throws ArgumentError reinterpret(Int32, Av) # ("cannot reinterpret `Vector{Int32}` as `Int32`, type `Vector{Int32}` is not a bits type")
+    @test_throws ArgumentError("cannot reinterpret a zero-dimensional `Int64` array to `Int32` which is of a different size") reinterpret(Int32, reshape([Int64(0)]))
+    @test_throws ArgumentError("cannot reinterpret a zero-dimensional `Int32` array to `Int64` which is of a different size") reinterpret(Int64, reshape([Int32(0)]))
+    @test_throws ArgumentError reinterpret(Tuple{Int,Int}, [1,2,3,4,5]) # ("""cannot reinterpret an `$Int` array to `Tuple{$Int, $Int}` whose first dimension has size `5`.
+                               # The resulting array would have non-integral first dimension.
+                               # """)
+    @test_throws ArgumentError("`reinterpret(reshape, Complex{Int64}, a)` where `eltype(a)` is Int64 requires that `axes(a, 1)` (got Base.OneTo(4)) be equal to 1:2 (from the ratio of element sizes)") reinterpret(reshape, Complex{Int64}, A)
+    @test_throws ArgumentError("`reinterpret(reshape, T, a)` requires that one of `sizeof(T)` (got 24) and `sizeof(eltype(a))` (got 16) be an integer multiple of the other") reinterpret(reshape, NTuple{3, Int64}, _B)
+    @test_throws ArgumentError reinterpret(reshape, Vector{Int64}, Ar) # ("cannot reinterpret `Int64` as `Vector{Int64}`, type `Vector{Int64}` is not a bits type")
+    @test_throws ArgumentError("cannot reinterpret a zero-dimensional `UInt8` array to `UInt16` which is of a larger size") reinterpret(reshape, UInt16, reshape([0x01]))
+
+    # getindex
+    _A = A
+    @test reinterpret(Complex{Int64}, _A) == [1 + 2im, 3 + 4im]
+    @test reinterpret(Float64, _A) == reinterpret.(Float64, A)
+    @test reinterpret(reshape, Float64, _A) == reinterpret.(Float64, A)
+
+    Ars = Ar
+    @test reinterpret(reshape, Complex{Int64}, Ar) == [1 + 2im, 3 + 4im]
+    @test reinterpret(reshape, Float64, Ar) == reinterpret.(Float64, Ars)
+
+    # setindex
+    A3 = collect(reshape(1:18, 2, 3, 3))
+    A3r = reinterpret(reshape, Complex{Int}, A3)
+    @test A3r[4] === A3r[1,2] === A3r[CartesianIndex(1, 2)] === 7+8im
+    A3r[2,3] = -8-15im
+    @test A3[1,2,3] == -8
+    @test A3[2,2,3] == -15
+    A3r[4] = 100+200im
+    @test A3[1,1,2] == 100
+    @test A3[2,1,2] == 200
+    A3r[CartesianIndex(1,2)] = 300+400im
+    @test A3[1,1,2] == 300
+    @test A3[2,1,2] == 400
+
+    # Test 0-dimensional Arrays
+    A = zeros(UInt32)
+    B = reinterpret(Int32,A)
+    Brs = reinterpret(reshape,Int32,A)
+    @test size(B) == size(Brs) == ()
+    @test axes(B) == axes(Brs) == ()
+    B[] = Int32(5)
+    @test B[] === Int32(5)
+    @test Brs[] === Int32(5)
+    @test A[] === UInt32(5)
+
+    # reductions
+    a = [(1,2,3), (4,5,6)]
+    ars = reinterpret(reshape, Int, a)
+    @test sum(ars) == 21
+    @test sum(ars; dims=1) == [6 15]
+    @test sum(ars; dims=2) == reshape([5,7,9], (3, 1))
+    @test sum(ars; dims=(1,2)) == reshape([21], (1, 1))
+    # also test large sizes for the pairwise algorithm
+    a = [(k,k+1,k+2) for k = 1:3:4000]
+    ars = reinterpret(reshape, Int, a)
+    @test sum(ars) == 8010003
+end
+
 include("iterators.jl")
 
 nothing