Improve type stability of cat_pullback

[ToucheSir]

Currently,

ChainRules.jl/src/rulesets/Base/array.jl

Lines 356 to 362 in 8c34f19

	index = ntuple(ndimsY) do d
	if d in cdims
	d > ndimsX ? (prev[d]+1) : (prev[d]+1:prev[d]+sizeX[d])
	else
	d > ndimsX ? 1 : (:)
	end
	end

will generate a Tuple{Union{UnitRange{Int64}, Colon}, ...}. When the number of input dims >=3, this appears to fail some union-splitting threshold for getindex on certain array types. Here's a MWE with FillArrays, which Zygote uses somewhat extensively:

using Test, ChainRules, ChainRulesCore, FillArrays

function grad(x1, x2, sens, dims::Val)
  y, back = rrule(cat, x1, x2; dims)
  back(sens)
end

# Infers
@inferred grad(ones(1, 1, 1), ones(1, 1, 1), ones(1, 2, 1), Val(2))
@inferred grad(Ones(1, 1), Ones(1, 1), ones(1, 2), Val(2))

# Does not infer
@inferred grad(Ones(1, 1, 1), Ones(1, 1, 1), Ones(1, 2, 1), Val(2))
@inferred grad(ones(1, 1, 1), ones(1, 1, 1), Ones(1, 2, 1), Val(2))

This PR resolves the instability by always creating a UnitRange. However, I am not sure if this is sufficiently semantically close to Colon to pass muster, so RFC :) .

[mcabbott]

This seems fine to me. Could you add something like this example as a test?

The rules for hcat and vcat have a similar branch, d > ndimsX ? 1 : (:). But perhaps that never matters.

[ToucheSir]

The rules for hcat and vcat have a similar branch, d > ndimsX ? 1 : (:). But perhaps that never matters.

It also doesn't matter for 1 and 2D inputs here, which is why I suspect union splitting is bailing out. Cthulhu shows 4 variants of getindex(::Fill, ...) generated for the 2D case, so I guess the threshold is somewhere between 2^2 and 2^3?