Guard against invalid ranks

Author: christiangnrd

lib/mps/ndarray.jl

@@ -7,6 +7,8 @@ export MPSNDArrayDescriptor
 # @objcwrapper immutable=false MPSNDArrayDescriptor <: NSObject
 
 function MPSNDArrayDescriptor(dataType::DataType, dimensionCount, dimensionSizes::Ptr)
+    1 <= dimensionCount <= 16 || throw(ArgumentError("`dimensionCount` must be between 1 and 16 inclusive"))
+
     desc = @objc [MPSNDArrayDescriptor descriptorWithDataType:dataType::MPSDataType
                                       dimensionCount:dimensionCount::NSUInteger
                                       dimensionSizes:dimensionSizes::Ptr{NSUInteger}]::id{MPSNDArrayDescriptor}

lib/mpsgraphs/tensor.jl

@@ -63,8 +63,9 @@ function MPSGraphTensorData(matrix::MPSMatrix)
     return tensor
 end
 
-# rank must be between 1 and 16 inclusive
 function MPSGraphTensorData(matrix::MPSMatrix, rank)
+    1 <= rank <= 16 || throw(ArgumentError("`rank` must be between 1 and 16 inclusive"))
+
     obj = @objc [MPSGraphTensorData alloc]::id{MPSGraphTensorData}
     tensor = MPSGraphTensorData(obj)
     finalizer(release, tensor)
@@ -81,8 +82,9 @@ function MPSGraphTensorData(vector::MPSVector)
     return tensor
 end
 
-# rank must be between 1 and 16 inclusive
 function MPSGraphTensorData(vector::MPSVector, rank)
+    1 <= rank <= 16 || throw(ArgumentError("`rank` must be between 1 and 16 inclusive"))
+
     obj = @objc [MPSGraphTensorData alloc]::id{MPSGraphTensorData}
     tensor = MPSGraphTensorData(obj)
     finalizer(release, tensor)

test/mps/ndarray.jl

@@ -29,6 +29,8 @@ end
     desc2.numberOfDimensions = 6
     @test desc2.numberOfDimensions == 6
 
+    @test_throws ArgumentError MPSNDArrayDescriptor(Float32, ones(Int, 17))
+
     @static if Metal.macos_version() >= v"15"
         @test desc1.preferPackedRows == false