MCPcopy Create free account

hub / github.com/ml-explore/mlx / functions

Functions4,946 in github.com/ml-explore/mlx

Methodtest_wired_memory
(self)
python/tests/test_memory.py:52
Methodtest_wrap_compiled
(self)
python/tests/test_compile.py:1102
Methodthread_count
mlx/3rdparty/pocketfft.h:498
Methodthread_fn
mlx/scheduler.h:37
Methodthread_fn
mlx/backend/cuda/worker.cpp:59
Functionthread_map
mlx/3rdparty/pocketfft.h:525
Methodthread_pool
mlx/3rdparty/pocketfft.h:720
Functionthread_reduce
mlx/backend/metal/kernels/reduction/reduce_row.h:165
Functionthreadgroup_reduce
mlx/backend/metal/kernels/reduction/reduce_row.h:129
Functiontile_matmad_nax
mlx/backend/metal/kernels/steel/attn/nax.h:825
Functiontime_fn
benchmarks/cpp/time_utils.h:26
Functionto_bnns_dtype<bfloat16_t>
mlx/backend/cpu/gemms/bnns.cpp:24
Functionto_bnns_dtype<float16_t>
mlx/backend/cpu/gemms/bnns.cpp:19
Functionto_bnns_dtype<float>
mlx/backend/cpu/gemms/bnns.cpp:15
Functionto_list
python/src/convert.cpp:205
Methodto_quantized
Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer.
python/mlx/nn/layers/embedding.py:43
Functiontolist
python/src/convert.cpp:219
Functiontopk
(axis, x)
benchmarks/python/comparative/bench_mlx.py:328
Functiontopk
(axis, x)
benchmarks/python/comparative/bench_torch.py:289
Methodtorch_fn
(*args)
python/tests/test_bf16.py:90
Methodtrainable_parameter_filter
(module, key, value)
python/mlx/nn/layers/base.py:239
Methodtraining
Boolean indicating if the model is in training mode.
python/mlx/nn/layers/base.py:67
Methodtransform
(x: State)
python/tests/test_autograd.py:914
Methodtransform
(x: State)
python/tests/test_vmap.py:781
Methodtransform
(x: State)
python/tests/test_compile.py:1242
Methodtransform_tuple
(t)
python/tests/test_autograd.py:917
Methodtransform_tuple
(t)
python/tests/test_vmap.py:784
Methodtransform_tuple
(t)
python/tests/test_compile.py:1245
Methodtransform_vector
(t)
python/tests/test_autograd.py:920
Methodtransform_vector
(t)
python/tests/test_vmap.py:787
Methodtransform_vector
(t)
python/tests/test_compile.py:1248
Functiontranspose_add
(a, b)
benchmarks/python/single_ops.py:18
Functiontri_inv
mlx/backend/cpu/inverse.cpp:72
Methodtrim
mlx/backend/cuda/lru_cache.h:122
Functiontriplet_loss
r""" Computes the triplet loss for a set of anchor, positive, and negative samples. Margin is represented with alpha in the math section.
python/mlx/nn/losses.py:386
Methodtwsize
mlx/3rdparty/pocketfft.h:1498
Methodtwsize
mlx/3rdparty/pocketfft.h:2297
Functionuint16_to_bfloat16
mlx/backend/metal/kernels/bf16.h:14
Functionunary_fused_0
No fusion
tests/compile_tests.cpp:239
Functionunary_fused_1_copy
tests/compile_tests.cpp:248
Functionunary_fused_1_diff
tests/compile_tests.cpp:252
Functionunary_fused_2
Output into un-compilable primitive
tests/compile_tests.cpp:257
Functionunary_fused_3
Input from un-compilable primitive
tests/compile_tests.cpp:262
Functionunary_g
mlx/backend/metal/kernels/unary.h:46
Functionunary_op
mlx/backend/cpu/unary.h:13
Functionunary_v
mlx/backend/metal/kernels/unary.h:4
Functionunary_v2
mlx/backend/metal/kernels/unary.h:22
Functionunary_with_two_outputs
Uncompilable input outside fused tape
tests/compile_tests.cpp:447
Functionunbatch_vjp_first
()
benchmarks/python/batch_matmul_bench.py:41
Functionunbatch_vjp_second
()
benchmarks/python/batch_matmul_bench.py:46
Methodunfreeze
Wrap unfreeze so that we unfreeze any layers we might contain but our parameters will remain frozen.
python/mlx/nn/layers/distributed.py:560
Methodunsafe_weak_copy
mlx/array.cpp:60
Methodupdate
Apply the gradients to the parameters of the model and update the model with the new parameters. Args: model (mlx.nn.Modu
python/mlx/optimizers/optimizers.py:20
Methodupdate
()
python/tests/test_optimizers.py:485
Methodupdate
mlx/backend/no_gpu/fence.cpp:39
Methodupdate
mlx/backend/metal/fence.cpp:100
Methodupdate
Non atomic update
mlx/backend/metal/kernels/reduction/ops.h:67
Methodupdate
Non atomic update
mlx/backend/metal/kernels/reduction/ops.h:108
Methodupdate_state
(params, state)
python/mlx/optimizers/optimizers.py:53
Methoduse_fallback
mlx/backend/no_gpu/primitives.cpp:26
Methoduse_fallback
mlx/backend/metal/rope.cpp:10
Methoduse_fallback
mlx/backend/metal/scaled_dot_product_attention.cpp:588
Methoduse_fallback
mlx/backend/metal/normalization.cpp:13
Methoduse_fallback
mlx/backend/cuda/scaled_dot_product_attention.cpp:552
Methodvalid_child_filter
(module, key, value)
python/mlx/nn/layers/base.py:231
Functionvalidate_subtrees
python/src/trees.cpp:6
Functionvalue_and_grad
mlx/transforms.cpp:666
Methodvalue_type
mlx/backend/metal/kernels/steel/utils/integral_constant.h:23
Functionvec<IdxT, 2> elem_to_loc_2_nd
mlx/backend/metal/kernels/utils.h:150
Functionvec<IdxT, 3> elem_to_loc_3_nd
mlx/backend/metal/kernels/utils.h:173
Methodvjp
mlx/primitives.cpp:189
Methodvjp
mlx/compile.cpp:152
Methodvjp
mlx/distributed/primitives.cpp:46
Methodvmap
mlx/primitives.cpp:201
Methodvmap
mlx/compile.cpp:167
Methodvmap
mlx/distributed/primitives.cpp:12
Methodvmap
Vectorize primitive along given axis */
examples/extensions/axpby/axpby.cpp:294
Methodwait
mlx/backend/no_gpu/event.cpp:22
Methodwait
mlx/backend/no_gpu/fence.cpp:23
Methodwait
mlx/backend/metal/event.cpp:23
Methodwait
mlx/backend/metal/fence.cpp:52
Methodwith_padded_channels
mlx/backend/metal/kernels/steel/conv/params.h:23
Methodworker
mlx/distributed/ring/ring.cpp:202
Functionwrapped_checkpointed_fn
(*args, **kwargs)
python/mlx/nn/utils.py:68
Functionwrapped_value_grad_fn
(*args, **kwargs)
python/mlx/nn/utils.py:34
Methodwrapper
()
python/tests/test_compile.py:1108
Functionwrite_padded
mlx/backend/metal/kernels/fft/readwrite.h:163
Functionwrite_safe
mlx/backend/metal/kernels/scan.h:213
Functionwrite_strided
mlx/backend/metal/kernels/fft/readwrite.h:210
Functionwrite_unsafe
mlx/backend/metal/kernels/scan.h:200
Methodzeros_fun
(x)
python/tests/test_compile.py:489
Function~Allocator
mlx/allocator.h:49
Function~ArrayDesc
mlx/array.h:504
Method~ArrayDesc
mlx/array.cpp:275
Method~BufferCache
mlx/backend/common/buffer_cache.h:23
Method~CaptureContext
mlx/backend/cuda/device.cpp:80
Method~CommandEncoder
mlx/backend/metal/device.cpp:281
Method~CommandEncoder
mlx/backend/cuda/device.cpp:215
Method~ConcurrentContext
mlx/backend/metal/device.h:36
Method~ConcurrentContext
mlx/backend/cuda/device.cpp:98
← previousnext →4,801–4,900 of 4,946, ranked by callers