Scaled dot product attention backward does not work with MacOS

snippet: ``` let keys = tokeys(c).reshaped(.NHWC(b, t, h, k)).identity() var queries = ( toqueries(x)).reshaped(.NHWC(b, hw, h, k)).identity().identity() var values = tovalues(c).reshaped(.NHWC(b, t, h, k)) let scaledDotProductAttention = ScaledDotProductAttention( scale: 1.0 / Float(k).squareRoot(), multiHeadOutputProjectionFused: true) var out = scaledDotProductAttention(queries, keys, values).reshaped([b, hw, h * k]) ``` error when MFA is enabled : ``` loc("mps_matmul_1"("(mpsFileLoc): ../Files/MPSGraphUtilities.mm":39:0)): error: incompatible dimensions loc("mps_matmul_1"("(mpsFileLoc): ..Files/MPSGraphUtilities.mm":39:0)): error: invalid shape ``` error when mfa is disabled ``` Assertion failed: (isStaticMPSType(type)), function setStaticJITypeForValue, file MPSRuntime_Project.h, line 501. ``` fwd works fine but in backward i get error