Functions — github.com/ml-explore/mlx

Methodrank

mlx/distributed/jaccl/lib/jaccl/mesh.h:28

Methodrank

mlx/distributed/nccl/nccl.cpp:324

Methodraw_ptr

mlx/backend/no_gpu/allocator.cpp:72

Methodraw_ptr

mlx/backend/metal/allocator.cpp:23

Functionrecv

mlx/distributed/distributed.cpp:41

Methodrecv

mlx/distributed/utils.cpp:150

Methodrecv

mlx/distributed/ring/ring.cpp:149

Methodrecv

mlx/distributed/ring/ring.cpp:291

Methodrecv

mlx/distributed/mpi/mpi.cpp:462

Methodrecv

mlx/distributed/jaccl/jaccl.cpp:138

Methodrecv

mlx/distributed/jaccl/lib/jaccl/ring.cpp:181

Methodrecv

mlx/distributed/jaccl/lib/jaccl/ring_impl.h:473

Methodrecv

mlx/distributed/jaccl/lib/jaccl/mesh.cpp:183

Methodrecv

mlx/distributed/jaccl/lib/jaccl/mesh_impl.h:325

Methodrecv

mlx/distributed/nccl/nccl.cpp:364

Methodrecv_from

mlx/distributed/jaccl/lib/jaccl/ring_impl.h:551

Methodrecv_from

mlx/distributed/jaccl/lib/jaccl/mesh_impl.h:378

Methodrecv_impl

mlx/distributed/ring/ring.cpp:181

Functionreduce_dispatch_and_or

mlx/backend/cpu/reduce.cpp:401

Functionreduce_dispatch_min_max

mlx/backend/cpu/reduce.cpp:435

Functionreduce_dispatch_sum_prod

mlx/backend/cpu/reduce.cpp:414

Methodreduce_scatter_impl

mlx/distributed/nccl/nccl.cpp:409

Functionreduction

(op, axis, x)

benchmarks/python/comparative/bench_mlx.py:140

Functionreduction

(op, axis, x)

benchmarks/python/comparative/bench_torch.py:98

Functionreduction_op

mlx/backend/cpu/reduce.cpp:139

Methodref_block_masked_mm

( a, b, block_size, out_mask=None, lhs_mask=None, rhs_mask=None )

python/tests/test_blas.py:839

Functionregular_attention

(x)

benchmarks/python/einsum_bench.py:56

Functionreleasebuffer

python/src/buffer.h:122

Functionrelu

r"""Applies the Rectified Linear Unit. Simply ``mx.maximum(x, 0)``.

python/mlx/nn/layers/activations.py:30

Functionrelu

(x)

benchmarks/python/comparative/bench_mlx.py:171

Functionrelu

(x)

benchmarks/python/comparative/bench_torch.py:132

Functionrelu2

r"""Applies the ReLU² activation function. Applies :math:`\max(0, x)^2` element wise.

python/mlx/nn/layers/activations.py:39

Functionrelu6

r"""Applies the Rectified Linear Unit 6. Applies :math:`\min(\max(x, 0), 6)` element wise.

python/mlx/nn/layers/activations.py:48

Functionrelu6

(x)

benchmarks/python/comparative/bench_mlx.py:220

Functionrelu6

(x)

benchmarks/python/comparative/bench_torch.py:164

Methodremove_from_list

mlx/backend/common/buffer_cache.h:130

Functionrepr_array

(self, x, maxlevel)

python/mlx/_reprlib_fix.py:9

Methodreset

mlx/backend/common/utils.h:162

Functionreset_peak_memory

mlx/backend/no_gpu/allocator.cpp:112

Functionreset_peak_memory

mlx/backend/metal/allocator.cpp:258

Functionreshape_transposed

()

benchmarks/python/single_ops.py:108

Methodresize

mlx/small_vector.h:403

Methodresolve_stream

mlx/export.cpp:477

Functionresult_type

mlx/utils.cpp:128

Methodretain_graph

mlx/transforms_impl.h:59

Methodreverse_sequence

(xs, axis=0)

python/tests/test_conv.py:932

Methodrfftp

mlx/3rdparty/pocketfft.h:2346

Functionrms_norm

mlx/fast.cpp:53

Functionrms_norm_loop

(g, x, w)

benchmarks/python/rms_norm_bench.py:29

Functionrope

(x)

benchmarks/python/comparative/bench_torch.py:244

Functionrope_mat

(x)

benchmarks/python/rope_bench.py:26

Functionrope_vec

(x)

benchmarks/python/rope_bench.py:15

Methodrow_bin_op

mlx/backend/metal/kernels/steel/gemm/nax.h:374

Methodrow_bin_op

mlx/backend/metal/kernels/steel/gemm/nax.h:653

Methodrow_bin_op

mlx/backend/metal/kernels/steel/attn/nax.h:374

Methodrow_bin_op

mlx/backend/metal/kernels/steel/attn/nax.h:653

Methodrow_bin_op

mlx/backend/metal/kernels/steel/attn/mma.h:215

Methodrow_bin_op

mlx/backend/metal/kernels/steel/attn/mma.h:307

Methodrow_reduce

mlx/backend/metal/kernels/steel/gemm/nax.h:354

Methodrow_reduce

mlx/backend/metal/kernels/steel/gemm/nax.h:640

Methodrow_reduce

mlx/backend/metal/kernels/steel/attn/nax.h:354

Methodrow_reduce

mlx/backend/metal/kernels/steel/attn/nax.h:640

Methodrow_reduce

mlx/backend/metal/kernels/steel/attn/mma.h:200

Methodrow_reduce

mlx/backend/metal/kernels/steel/attn/mma.h:295

Functionrow_reduce_looped

mlx/backend/metal/kernels/reduction/reduce_row.h:314

Functionrow_reduce_simple

mlx/backend/metal/kernels/reduction/reduce_row.h:266

Functionrow_reduce_small

mlx/backend/metal/kernels/reduction/reduce_row.h:199

Methodrun

mlx/backend/cpu/eig.cpp:65

Methodrun

mlx/backend/cpu/svd.cpp:69

Methodrun

mlx/backend/cpu/eigh.cpp:54

Methodrun

mlx/backend/metal/kernels/gemv_masked.h:125

Methodrun

mlx/backend/metal/kernels/gemv_masked.h:404

MethodrunTests

(self)

python/tests/mlx_tests.py:57

Methodrun_batched

mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp:9

Methodrun_conv1D

( N, C, O, iH, kH, stride,

python/tests/test_conv.py:60

Methodrun_conv1D_grad

( N, C, O, iH, kH, stride,

python/tests/test_conv.py:169

Methodrun_conv2D

( N, C, O, idim, kdim, stride,

python/tests/test_conv.py:275

Methodrun_conv2D_grad

( N, C, O, idim, kdim, stride,

python/tests/test_conv.py:375

Methodrun_conv3D

( N, C, O, idim, kdim, stride,

python/tests/test_conv.py:485

Methodrun_conv3D_grad

( N, C, O, idim, kdim, stride,

python/tests/test_conv.py:569

Methodrun_conv_transpose1D_grad

( N, C, O, iH, kH, stride,

python/tests/test_conv_transpose.py:132

Methodrun_conv_transpose2D

( N, C, O, idim, kdim, stride,

python/tests/test_conv_transpose.py:230

Methodrun_conv_transpose2D_grad

( N, C, O, idim, kdim, stride,

python/tests/test_conv_transpose.py:320

Methodrun_conv_transpose3D

( N, C, O, idim, kdim, stride,

python/tests/test_conv_transpose.py:417

Methodrun_conv_transpose3D_grad

( N, C, O, idim, kdim, stride,

python/tests/test_conv_transpose.py:495

Methodrun_conv_transpose_1D

( N, C, O, iH, kH, stride,

python/tests/test_conv_transpose.py:23

Methodrun_conv_transpose_1d_output_padding

( N, C, O, iH, kH, stride, padding, output_padding, dtype="float32", atol=1e-5 )

python/tests/test_conv_transpose.py:601

Methodrun_conv_transpose_2d_output_padding

( N, C, O, idim, kdim, stride,

python/tests/test_conv_transpose.py:657

Methodrun_conv_transpose_3d_output_padding

( N, C, O, idim, kdim, stride,

python/tests/test_conv_transpose.py:732

Methodrun_test

(a, b, block_size, out_mask, a_mask, b_mask, cotan)

python/tests/test_blas.py:879

Methodrun_test_mask_vjp

(a, b, block_size, out_mask, a_mask, b_mask, cotan)

python/tests/test_blas.py:895

Methodrun_upsample

( N, C, idim, scale_factor, mode, alig

python/tests/test_upsample.py:22

Functionsafe_divup

mlx/backend/metal/reduce.cpp:192

Methodsanity_check

mlx/3rdparty/pocketfft.h:464

Functionsave_gguf

mlx/io/no_gguf.cpp:12

Functionsave_safetensors

mlx/io/no_safetensors.cpp:19

Functionscalar_mult

(x)

benchmarks/python/comparative/bench_mlx.py:248

Functionscalar_mult

(x)

benchmarks/python/comparative/bench_torch.py:204

Functionscan_dispatch

mlx/backend/cpu/scan.cpp:194

Functionscan_op

mlx/backend/cpu/scan.cpp:157

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