apl: Phase 3 squad ⌷ indexing — 66/66 tests green
Some checks failed
Test, Build, and Deploy / test-build-deploy (push) Has been cancelled
Some checks failed
Test, Build, and Deploy / test-build-deploy (push) Has been cancelled
Add apl-squad: scalar index into vector, fully-specified multi-dim index, partial index returning sub-array slice. 7 new tests. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -661,3 +661,30 @@
|
||||
(make-array
|
||||
(cons (+ (first a-s) (first b-s)) (rest a-s))
|
||||
(append a-r b-r)))))
|
||||
|
||||
(define
|
||||
apl-squad
|
||||
(fn
|
||||
(idx-arr data-arr)
|
||||
(let
|
||||
((shape (get data-arr :shape))
|
||||
(ravel (get data-arr :ravel))
|
||||
(strides (apl-strides (get data-arr :shape))))
|
||||
(let
|
||||
((idxs (if (scalar? idx-arr) (list (disclose idx-arr)) (get idx-arr :ravel))))
|
||||
(let
|
||||
((k (len idxs)) (rank (len shape)))
|
||||
(let
|
||||
((adj (map (fn (i) (- i apl-io)) idxs)))
|
||||
(if
|
||||
(= k rank)
|
||||
(apl-scalar (nth ravel (apl-multi->flat adj strides)))
|
||||
(let
|
||||
((remaining-shape (drop shape k))
|
||||
(start (apl-multi->flat adj (take strides k)))
|
||||
(slice-size (reduce * 1 (drop shape k))))
|
||||
(make-array
|
||||
remaining-shape
|
||||
(map
|
||||
(fn (j) (nth ravel (+ start j)))
|
||||
(range 0 slice-size)))))))))))
|
||||
|
||||
Reference in New Issue
Block a user