Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Indices which run over a shorter range than axes(A,d) #56

Open
prittjam opened this issue Oct 6, 2022 · 9 comments
Open

Indices which run over a shorter range than axes(A,d) #56

prittjam opened this issue Oct 6, 2022 · 9 comments

Comments

@prittjam
Copy link

prittjam commented Oct 6, 2022

why doesn't this work?

P = rand(3,4,100)
@cast H[k][i,j] := P[i,j,k] (i in 1:3, j in 1:3, k in 1:100)

It fails with

DimensionMismatch("range of index j must agree")

Probably i am missing something basic.
Thanks.

@mcabbott
Copy link
Owner

mcabbott commented Oct 6, 2022

I see the temptation but this isn't supported. At present, every index of every array runs the full range. Thus in every operation combining two arrays, the dimensions sharing an index must agree.

This rule comes from broadcasting, where things like [1,2,3] .+ [1,2,3,4] just fail, rather than trimming one range. The only time at which things like i in 1:3 are needed is when reshaping one dimension to more than one, e.g. @cast m[i,j] := rand(6)[(i,j)] could describe several different shapes of matrix.

It's not crazy to imagine other rules, but they will add complexity both to the interface and to the implementation. It could be that j in 1:3 should take priority. (IIRC at present j in 1:3 is no different to having an array this size, i.e. there is no special priority.) Or it could be an intersection of all constraints, possibly including that a[i] + b[i] should run over shared elements. (In which case many size mistakes won't give errors.)

Besides complexity within the macro, one issue with allowing this would be that axes(P,2) != 1:3 is known only at runtime, and so the code it generates would probably always have to make something like view(P, :, 1:3, :). That's unfortunately likely to cause problems with GPU arrays, which don't like stacked wrappers.

@mcabbott
Copy link
Owner

mcabbott commented Oct 6, 2022

You can of course write @cast H[k][i,j] := view(P, :, 1:3, :)[i,j,k]. A feature I don't promise to implement would be to make the macro produce this when asked, something like @cast H[k][i,j] := P[i, (j in 1:3), k], without changing the overall rule, nor wrapping every array in a view.

Maybe also worth noting that Tullio has a slightly different rule for indices. By default it's the same, but as soon as an index appears with a shift like A[i+1], then this i runs over the intersection of all ranges. (TensorCast never allows such shifts.) Xref mcabbott/Tullio.jl#89

@prittjam
Copy link
Author

prittjam commented Oct 6, 2022

Thanks for the response! I think this is a common use case for casting, though: to subindex and slice at the same time. It may not be always possible to create a view because, e.g., the right-hand side may be of the form of a vector of matrices from a prior cast, e.g.,

@cast H[k][i,j] := P[k][i,j]

and here I don't see an elegant alternative. Here P[k][i,j] is 3x4 and we want H[k][i,j] to be 3x3.

@mcabbott
Copy link
Owner

mcabbott commented Oct 6, 2022

Indeed. And mixing a hand-written view with the macro doesn't go so well... ideally this should be smart enough not to assume that dolphin is a scalar, but it isn't... and still looks like a hack:

julia> vm = [rand(3,4) for _ in 1:5];

julia> @cast A[i,j,k] := view(vm[k], :, 1:2)[i,j];
ERROR: DimensionMismatch: arrays could not be broadcast to a common size; got a dimension with lengths 5 and 2

julia> @pretty @cast A[i,j,k] := view(vm[k], :, 1:2)[i,j];
begin
    @boundscheck vm isa Tuple || (ndims(vm) == 1 || throw(ArgumentError("expected a vector or tuple vm[k]")))
    local dolphin = 1:2
    local pig = @__dot__(view(vm, :, dolphin))  # @. hits the view too
    local armadillo = lazystack(pig)
    A = armadillo
end

julia> @cast A[i,j,k] |= view(vm[k], :, Ref(1:2))[i,j];  # this works but it's weird, really not encouraged!

julia> summary(A)
"3×2×5 Array{Float64, 3}"

@prittjam
Copy link
Author

prittjam commented Oct 6, 2022

I think this is a common enough use case where sane notation is needed. E.g., if I want to put these arrays into a struct array, working with matrices of the from P[k][i,j] is the most natural.

@mcabbott mcabbott changed the title subindices Indices which run over a shorter range than axes(A,d) Oct 6, 2022
@mcabbott
Copy link
Owner

mcabbott commented Oct 6, 2022

Some possible notations might be:

@cast A[i,j,k] |= vm[k][i, (j in 1:2)];

@cast A[i,j,k] |= vm[k][i, view(j)]  j in 1:2;  # maybe easier

If in the mood you can try to see where in macro.jl this needs to happen. It's all recursive and I forget how it works repeatedly...

It would also require a bit of thought about how this intersects with other features like reversed axes and (i,j)... or just to make sure all weird combinations are forbidden. Presumably @cast A[i, view(j)] = ... would be allowed, in-place, but not with :=, new array.

@prittjam
Copy link
Author

prittjam commented Oct 6, 2022

You might overestimate my Julia proficiency a bit :), but I can take a look.

@prittjam
Copy link
Author

prittjam commented Oct 6, 2022

I think some of the use cases for tensor cast would not be needed if eachslice accepted tuples for dims, i.e., eachslice(P,dims=(2,3)). Is there some technical reason why this feature is not there?

@mcabbott
Copy link
Owner

mcabbott commented Oct 6, 2022

Julia 1.9 will have this. And ideally JuliaLang/Compat.jl#663 would make this "official" version available before then.

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

No branches or pull requests

2 participants