mirror of
https://github.com/ml-explore/mlx.git
synced 2025-06-24 09:21:16 +08:00
11 lines
250 B
Python
11 lines
250 B
Python
|
import mlx.core as mx
|
||
|
|
from mlx_sample_extensions import axpby
|
||
|
|
|
||
|
|
a = mx.ones((3, 4))
|
||
|
|
b = mx.ones((3, 4))
|
||
|
|
c = axpby(a, b, 4.0, 2.0, stream=mx.cpu)
|
||
|
|
|
||
|
|
print(f"c shape: {c.shape}")
|
||
|
|
print(f"c dtype: {c.dtype}")
|
||
|
|
print(f"c correct: {mx.all(c == 6.0).item()}")
|