Conversation
|
Sorry for the churn -- I had to rename a couple times before circleci was happy with the branch name |
| public func transposed(axis: Int, stream: StreamOrDevice = .default) -> MLXArray { | ||
| var result = mlx_array_new() | ||
| mlx_transpose(&result, ctx, [axis.int32], 1, stream.ctx) | ||
| mlx_transpose_axes(&result, ctx, [axis.int32], 1, stream.ctx) |
There was a problem hiding this comment.
@andresy not a problem, but should we have a mlx_transpose_axis? This is the only one I found where we have API on the swift side that didn't have the matching mlx-c piece.
|
Note: tests will fail until tags are cut and we can update the Cmakelists |
| /// - array: input array | ||
| /// - stream: stream or device to evaluate on | ||
| /// - Returns: The `S` matrix | ||
| public static func svd(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { |
There was a problem hiding this comment.
How does this work with the svd above? Can you overload based on return type? In python we use the compute_uv parameter to determine what to return.
There was a problem hiding this comment.
Yeah, that is a good question. In swift we have to have two methods since we have different return types. We could return an enum for a union type, but that is not very convenient to use. You can overload based on return type and it will pick the right one, but I think we may want a hint here (disfavoredOverload).
Anyway this works:
let a = MLXRandom.uniform(0 ..< 1, [10, 10])
let (u, s, v) = MLXLinalg.svd(a)
let s2 = MLXLinalg.svd(a)| /// | ||
| /// * `B`: The batch size. | ||
| /// * `N_q`: The number of query heads. | ||
| /// * `N_kv`: The number of key and value heads. |
There was a problem hiding this comment.
I don't think it's a requirement that N_k == N_v. The function should still work if they are different.
There was a problem hiding this comment.
🤷
I copied it from the python docs:
| /// - values: values with shape `[B, N_kv, T_kv, D]` | ||
| /// - scale: scale for queries, typically `1 / sqrt(q.dim(-1))` | ||
| /// - mask: mask array | ||
| /// - memoryEfficientThreshold: unused |
There was a problem hiding this comment.
I thought we got rid of that parameter 🤔 , are you keeping it for back-compat?
| /// - mask: mask array | ||
| /// - memoryEfficientThreshold: unused | ||
| /// - stream: stream to evaluate on | ||
| public static func scaledDotProductAttention( |
There was a problem hiding this comment.
Should this function be marked deprecated? It seems like the one below is the one we want people to use.
There was a problem hiding this comment.
No, it is to handle this variant:
mask (Union[None, str, array], optional): A causal, boolean or additive
mask to apply to the query-key scores. The mask can have at most 4
dimensions and must be broadcast-compatible with the shape
``[B, N, T_q, T_kv]``. If an additive mask is given its type must
promote to the promoted type of ``q``, ``k``, and ``v``.
awni
left a comment
There was a problem hiding this comment.
LGTM! Left a few comments, please take a look and merge when ready.
- fixes #211 (moves metalKernel API back to what it was) - add improved swift error handling - add function import/export - adopt mlx v0.25.1 and mlx-c v0.2.0 - fixes #211 (moves metalKernel API back to what it was) - use an evalLock for eval, asyncEval and Stream creation - add withError and withErrorHandler error handling - add function import/export: https://ml-explore.github.io/mlx/build/html/usage/export.html
adopt mlx v0.25.1 and mlx-c v0.2.0
add saveToData and load from data (safetensors) -- fix Feature to encode MLXArrays to .safetensors binary data in-memory #214
fixes metalKernel() should recreate the kernel object each time #211 (moves metalKernel API back to what it was)
add withError and withErrorHandler error handling
add function import/export: https://ml-explore.github.io/mlx/build/html/usage/export.html
use an evalLock for eval, asyncEval and Stream creation