MCPcopy Create free account

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

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

⨍ Functions 4,946◇ Types & classes 770↳ Endpoints 5
↓ 1 callersMethod_approximate_exp_moving_avg
(self, exp_avg_sq_row, exp_avg_sq_col)
python/mlx/optimizers/optimizers.py:795
↓ 1 callersFunction_bs_qmm_dispatch
mlx/backend/cpu/quantized.cpp:749
↓ 1 callersMethod_calc_stats
Calculate the mean and variance of the input tensor across the batch and spatial dimensions. Args: x (array): In
python/mlx/nn/layers/normalization.py:318
↓ 1 callersMethod_compute_learning_rate
(self, step, parameter_rms)
python/mlx/optimizers/optimizers.py:782
↓ 1 callersFunction_drop_dim
(shape, axis)
python/mlx/nn/losses.py:73
↓ 1 callersMethod_extra_repr
(self)
python/mlx/nn/layers/base.py:84
↓ 1 callersFunction_nearest_indices
(N, scale, dim, ndims)
python/mlx/nn/layers/upsample.py:27
↓ 1 callersFunction_non_overlapping_sliding_windows
(x, shape, window_shape)
python/mlx/nn/layers/pooling.py:23
↓ 1 callersFunction_qmm_dispatch
mlx/backend/cpu/quantized.cpp:421
↓ 1 callersFunction_sdpa
(q, k, v)
benchmarks/python/sdpa_vector_bench.py:24
↓ 1 callersMethod_set_training_mode
(self, mode: bool)
python/mlx/nn/layers/base.py:569
↓ 1 callersFunction_sliding_windows
(x, window_shape, window_strides)
python/mlx/nn/layers/pooling.py:39
↓ 1 callersFunction_unwrap
(model, value_key, value, filter_fn, map_fn, is_leaf_fn)
python/mlx/nn/layers/base.py:649
↓ 1 callersFunction_update_modules
(dst, modules, strict)
python/mlx/nn/layers/base.py:619
↓ 1 callersMethod_zeropower_via_newtonschulz5
(self, X, steps: int)
python/mlx/optimizers/optimizers.py:896
↓ 1 callersMethodadd_completed_handler
mlx/backend/cuda/device.cpp:220
↓ 1 callersMethodadd_task
mlx/backend/cuda/worker.cpp:31
↓ 1 callersFunctionaffine_dequantize
mlx/ops.cpp:4954
↓ 1 callersFunctionaffine_dequantize
mlx/backend/metal/kernels/quantized.h:2537
↓ 1 callersFunctionaffine_quantize
mlx/ops.cpp:4735
↓ 1 callersFunctionaffine_quantize
mlx/backend/metal/kernels/quantized.h:2432
↓ 1 callersMethodall_gather
mlx/distributed/distributed.cpp:75
↓ 1 callersMethodall_gather
mlx/distributed/jaccl/lib/jaccl/mesh.cpp:174
↓ 1 callersFunctionall_gather_benchmark
()
benchmarks/python/synchronize_bench.py:39
↓ 1 callersMethodall_max
mlx/distributed/distributed.cpp:88
↓ 1 callersMethodall_min
mlx/distributed/distributed.cpp:93
↓ 1 callersFunctionall_reduce_benchmark
()
benchmarks/python/synchronize_bench.py:23
↓ 1 callersFunctionall_reduce_dispatch
mlx/backend/metal/reduce.cpp:312
↓ 1 callersMethodall_sum
mlx/distributed/distributed.cpp:71
↓ 1 callersMethodapply
(dst, parameters)
python/mlx/nn/layers/base.py:325
↓ 1 callersMethodapply
mlx/backend/metal/kernels/steel/gemm/transforms.h:29
↓ 1 callersMethodapply
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h:32
↓ 1 callersMethodapply_single
Performs the SGD parameter update and stores :math:`v` in the optimizer state.
python/mlx/optimizers/optimizers.py:272
↓ 1 callersFunctionarange
mlx/backend/metal/kernels/arange.h:3
↓ 1 callersFunctionargmin
mlx/ops.cpp:2399
↓ 1 callersFunctionarray_basics
examples/cpp/tutorial.cpp:10
↓ 1 callersFunctionasync_eval
mlx/transforms.h:15
↓ 1 callersFunctionattention
(q, k, v, mask=None, w=None)
benchmarks/python/sdpa_vector_bench.py:23
↓ 1 callersFunctionautomatic_differentiation
examples/cpp/tutorial.cpp:81
↓ 1 callersFunctionbandwidth_gb
(runtime_ms, system_size)
benchmarks/python/fft_bench.py:14
↓ 1 callersMethodbarrier
mlx/backend/metal/device.cpp:372
↓ 1 callersFunctionbatch_tensordot
mlx/einsum.cpp:349
↓ 1 callersFunctionbench_case
(length, density, dtype)
benchmarks/python/masked_scatter.py:148
↓ 1 callersFunctionbench_gelu
()
benchmarks/python/compile_bench.py:11
↓ 1 callersFunctionbench_gemm
(M, N, K, dtype, rtol)
benchmarks/python/large_gemm_bench.py:56
↓ 1 callersFunctionbench_layernorm
()
benchmarks/python/compile_bench.py:52
↓ 1 callersFunctionbench_mlx
(steps: int = 20, shape=(10, 32, 32, 32, 3))
benchmarks/python/conv3d_train_bench_cpu.py:9
↓ 1 callersFunctionbench_mlx
(a, b, segments, warmup, iters)
benchmarks/python/segmented_mm_bench.py:52
↓ 1 callersFunctionbench_mlx
(steps: int = 20)
benchmarks/python/conv2d_train_bench_cpu.py:9
↓ 1 callersFunctionbench_mlx
(a, b)
benchmarks/python/large_gemm_bench.py:14
↓ 1 callersFunctionbench_mlx_loop
(a, b, segments, warmup, iters)
benchmarks/python/segmented_mm_bench.py:67
↓ 1 callersFunctionbench_shape
(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv3d_bench.py:50
↓ 1 callersFunctionbench_shape
(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv2d_transpose_bench_cpu.py:53
↓ 1 callersFunctionbench_shape
(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv_transpose_bench.py:57
↓ 1 callersFunctionbench_shape
(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv_bench.py:56
↓ 1 callersFunctionbench_shape
(N, iH, C, wH, O, strides, padding, np_dtype, groups)
benchmarks/python/conv1d_bench.py:56
↓ 1 callersFunctionbench_shape
(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv3d_bench_cpu.py:50
↓ 1 callersFunctionbench_shape
(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv_unaligned_bench.py:50
↓ 1 callersFunctionbench_shape
( B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, dtype, transpose=True, mask_in=None )
benchmarks/python/sdpa_bench.py:129
↓ 1 callersFunctionbench_shape
(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv3d_transpose_bench_cpu.py:54
↓ 1 callersFunctionbench_shape
(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype)
benchmarks/python/conv2d_bench_cpu.py:50
↓ 1 callersFunctionbench_shape
(B, M, N, K, np_dtype, transpose="nn")
benchmarks/python/blas/bench_gemm.py:109
↓ 1 callersFunctionbench_torch
(steps: int = 20, shape=(10, 3, 32, 32, 32))
benchmarks/python/conv3d_train_bench_cpu.py:73
↓ 1 callersFunctionbench_torch
(steps: int = 20)
benchmarks/python/conv2d_train_bench_cpu.py:73
↓ 1 callersFunctionbench_torch
(a, b)
benchmarks/python/large_gemm_bench.py:29
↓ 1 callersFunctionbench_with_in_len
(ax, in_vec_len, out_vector_lens, dtype, transpose)
benchmarks/python/blas/bench_gemv.py:135
↓ 1 callersFunctionbench_with_out_len
(ax, out_vec_len, in_vector_lens, dtype, transpose)
benchmarks/python/blas/bench_gemv.py:166
↓ 1 callersFunctionbenchmark_gather_mlx
(x_shape, idx_shape)
benchmarks/python/gather_bench.py:10
↓ 1 callersFunctionbenchmark_gather_torch
(x_shape, idx_shape, device)
benchmarks/python/gather_bench.py:21
↓ 1 callersFunctionbenchmark_scatter_mlx
(dst_shape, x_shape, idx_shapes)
benchmarks/python/scatter_bench.py:10
↓ 1 callersFunctionbenchmark_scatter_torch
(dst_shape, x_shape, idx_shapes, device)
benchmarks/python/scatter_bench.py:25
↓ 1 callersFunctionbenchmark_slice_update_mlx
(dst_shape, slice_shape, slice_range, dtype, iters=10)
benchmarks/python/slice_update_bench.py:10
↓ 1 callersFunctionbenchmark_slice_update_torch
( dst_shape, slice_shape, slice_range, device, dtype, iters=10 )
benchmarks/python/slice_update_bench.py:32
↓ 1 callersFunctionbfs_max_width
mlx/utils.h:149
↓ 1 callersFunctionbinary
(op, x, y)
benchmarks/python/comparative/bench_mlx.py:134
↓ 1 callersFunctionbinary_compile_outputs
tests/compile_tests.cpp:537
↓ 1 callersFunctionbinary_fused_1
Binary into unary
tests/compile_tests.cpp:338
↓ 1 callersFunctionbootstrap_unique_id
mlx/distributed/nccl/nccl.cpp:131
↓ 1 callersFunctionbuffer_format
python/src/buffer.h:20
↓ 1 callersFunctionbuild_case
(length, density, np_dtype, torch_dtype)
benchmarks/python/masked_scatter.py:116
↓ 1 callersFunctionbuild_conv_graph
mlx/backend/cuda/conv.cpp:93
↓ 1 callersFunctionbuild_kernel
mlx/backend/cpu/compiled.cpp:150
↓ 1 callersFunctionbuild_kernel
mlx/backend/cuda/custom_kernel.cpp:51
↓ 1 callersFunctionbuild_sdpa_backward_graph
mlx/backend/cuda/scaled_dot_product_attention.cpp:259
↓ 1 callersFunctionbuild_sdpa_graph
mlx/backend/cuda/scaled_dot_product_attention.cpp:208
↓ 1 callersFunctionbus_factor
Conversion from algorithm bandwidth to bus bandwidth for a ring reduce.
mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp:87
↓ 1 callersFunctionbytes_touched
(length, true_count, item_size)
benchmarks/python/masked_scatter.py:109
↓ 1 callersFunctioncan_dot
Assumes inputs have already have had repeats and single axis sums collapsed
mlx/einsum.cpp:335
↓ 1 callersFunctioncelu
r"""Applies the Continuously Differentiable Exponential Linear Unit. Applies :math:`\max(0, x) + \min(0, \alpha * (\exp(x / \alpha) - 1))` el
python/mlx/nn/layers/activations.py:125
↓ 1 callersFunctioncexpf
mlx/backend/metal/kernels/cexpf.h:83
↓ 1 callersFunctioncheck_correctness
(out_mx, out_pt, rtol, M, N, K)
benchmarks/python/large_gemm_bench.py:45
↓ 1 callersFunctioncheck_ssh_connections
(hosts, ignore_unreachable=False)
python/mlx/_distributed_utils/config.py:370
↓ 1 callersFunctioncheckpoint
mlx/transforms.cpp:1066
↓ 1 callersMethodclear_cache
mlx/backend/metal/allocator.cpp:167
↓ 1 callersMethodclear_library
mlx/backend/metal/device.cpp:715
↓ 1 callersMethodclose
python/src/export.cpp:84
↓ 1 callersFunctioncollapse_repeats
Collapse repeated subscripts and return the resulting array. The subscript is also updated in place. For example: - Given an input with shape (4, 4) a
mlx/einsum.cpp:424
↓ 1 callersMethodcommunication_stream
mlx/distributed/distributed.cpp:55
↓ 1 callersFunctioncompare
(args)
benchmarks/python/comparative/compare.py:24
↓ 1 callersFunctioncompare_mlx_dtypes
(args, dt1, dt2)
benchmarks/python/comparative/compare.py:31
← previousnext →1,201–1,300 of 4,946, ranked by callers