Split - #16
Split#16KhalidAl-otaibi wants to merge 2 commits into
Conversation
Not sure if you mean subscript or something else, but if you mean subscript we do have an impl for that |
| /// takes a shape for the original array, a shape for the subarray, | ||
| /// and an axis on which the original array would be divided | ||
| /// and returns an array of start and end indices for each subarray | ||
|
|
There was a problem hiding this comment.
For doc comments, let's use both starting capitalization and ending punctuation. also // helper function is missing an /
| /// and returns an array of start and end indices for each subarray | ||
|
|
||
| @usableFromInline | ||
| internal func _calculateSubarrayIndices(shape: [Int], subarrayShape: [Int], axis: Int) -> [SubarrayIndices] { |
There was a problem hiding this comment.
The only use of the_ prefix for a name is for public identifiers such as public protocol _ShapedArrayProtocol because they need to be public to work yet we want to signal to the user that they shouldn't be used directly. For non-public identifiers I think we should omit _, in this case that would be calculateSubarrayIndices.
However, i think we can improve on the name more too. Possibly
slicedSubArrayIndices(forShape shape: [Int], subarrayShape: [Int], alongAxis axis: Int) -> [SubarrayIndices]?
Lastly, and maybe this isn't possible with @usableFromInline, but if this is only used in 1 function I usually move the function declaration inside of the function that uses it. It adds some indentation but makes it clear at point of use that it isn't used elsewhere. Up to you though, nbd either way.
There was a problem hiding this comment.
or maybe splitSubArrayIndices to match the calling function..
| internal func _calculateSubarrayIndices(shape: [Int], subarrayShape: [Int], axis: Int) -> [SubarrayIndices] { | ||
| let numberOfSubarrays = shape[axis] / subarrayShape[axis] | ||
|
|
||
| var subarrays = [SubarrayIndices]() |
There was a problem hiding this comment.
Trying to think of a name that better captures this type, maybe var groupedSubarrayIndices? any suggestions?
| endIndices[axis] = startIndices[axis] + subarrayShape[axis] | ||
|
|
||
| subarrays.append(SubarrayIndices(start: startIndices, end: endIndices)) | ||
| } |
There was a problem hiding this comment.
this may look a bit nicer with a map
let subarrays = (0..<numberSubarrays).map { i in
var startIndices = [Int](repeating: 0, count: shape.count)
var endIndices = shape
startIndices[axis] = i * subarrayShape[axis]
endIndices[axis] = startIndices[axis] + subarrayShape[axis]
return SubarrayIndices(start: startIndices, end: endIndices)
}
| } | ||
|
|
||
| // helper function for subarray indices generation | ||
| /// takes a start and end index for each dimension and returns all the n-dim indices in between |
There was a problem hiding this comment.
Same comment as before about doc comments.
| /// Splits the ShapedArray into multiple subarrays along the given axis. | ||
| /// - Parameters: | ||
| /// - count: The number of subarrays to return. | ||
| /// - axis: The axis along which to split the ShapedArray. Negative values wrap around. |
There was a problem hiding this comment.
- axis seems misaligned with - count
| // helper function for subarray indices generation | ||
| /// takes a start and end index for each dimension and returns all the n-dim indices in between | ||
| @usableFromInline | ||
| internal func _generateIndices(start: [Int], end: [Int], currentIndex: [Int] = [], depth: Int = 0) -> [[Int]] { |
There was a problem hiding this comment.
Maybe
allIndices(fromStart start: [Int], toEnd end: [Int], currentIndices: [Int] = [], depth: Int = 0) -> [[Int]]
|
|
||
| for i in start[depth]..<end[depth] { | ||
| let newIndex = currentIndex + [i] | ||
| let subIndices = _generateIndices(start: start, end: end, currentIndex: newIndex, depth: depth + 1) |
There was a problem hiding this comment.
This isn't tail recursive so there's no tail-call optimization. Granted probably plenty to optimize in the future, but maybe add a comment here calling that out explicitly?
E.g. TODO: Possible tail-call optimization?
| } | ||
|
|
||
| // helper function for subarray indices generation | ||
| /// takes a shape, strides, and n-dim indices and returns the linear index |
| // helper function for subarray indices generation | ||
| /// takes a shape, strides, and n-dim indices and returns the linear index | ||
| @usableFromInline | ||
| internal func _calculateLinearIndex(shape: [Int], strides: [Int], indices: [Int]) -> Int { |
There was a problem hiding this comment.
Maybe
linearIndexFrom(shape: [Int], strides: [Int], indices: [Int]) -> Int?
other option may be bufferIndexFrom(shape:strides:indices)
| } | ||
| linearIndex += stride * index | ||
| } | ||
| return linearIndex |
| return linearIndex | ||
| } | ||
|
|
||
|
|
| return newArrays | ||
| } | ||
|
|
||
|
|
| /// the provided ShapedArrays. | ||
| /// | ||
|
|
||
|
|
There was a problem hiding this comment.
both newlines here should be deleted please, this breaks up the comment
| } | ||
|
|
||
| extension _ShapedArrayProtocol where Scalar: Comparable { | ||
| internal func _isLess(than other: Self) -> Bool { |
There was a problem hiding this comment.
I don't see this function used anywhere, is this something we need to keep?
|
|
||
| } | ||
|
|
||
| // test expandingShape |
There was a problem hiding this comment.
Comment here probably not needed since it's in the name of the test function
| return newArrays.map { ShapedArray(shape: newShape, scalars: $0) } | ||
| } | ||
|
|
||
| @inlinable |
There was a problem hiding this comment.
Some docs we can include:
/// Returns a shape-expanded `ShapedArray`, with a dimension of 1 inserted at the specified shape
/// indices.
| } | ||
|
|
||
| @inlinable | ||
| public func rankLifted() -> ShapedArray { |
There was a problem hiding this comment.
Possible docs for this:
/// Returns a rank-lifted `ShapedArray` with a leading dimension of 1.
rexmas
left a comment
rexmas
left a comment
There was a problem hiding this comment.
thank you for this! left some comments before approval, but nothing major
2 participants
Split function and other basic function such as expandingShape and squeezingShape. I think there is a much better implementation for the split function but this is what I came up with. It uses a few helper functions. These helper functions could also be used for subscription. In fact, if subscription was implemented, split could have been implemented very straightforwardly.