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
↓ 2 callersFunction_check_sharding
(sharding)
python/mlx/nn/layers/distributed.py:108
↓ 2 callersFunction_clip_grads_fsdp
(grads_slice, max_norm, group=None)
python/mlx/nn/utils.py:176
↓ 2 callersMethod_compute_rms
(self, inputs)
python/mlx/optimizers/optimizers.py:779
↓ 2 callersFunction_extract_info
(flat)
python/mlx/nn/utils.py:74
↓ 2 callersFunction_group_by_size
(keys, sizes, itemsize, communication_size)
python/mlx/nn/utils.py:82
↓ 2 callersFunction_interpolate
( x: mx.array, scale_factor: Tuple, indices_fn: Callable, align_corners: bool = False )
python/mlx/nn/layers/upsample.py:122
↓ 2 callersFunction_scaled_indices
(N, scale, align_corners, dim, ndims)
python/mlx/nn/layers/upsample.py:12
↓ 2 callersFunction_split
Equivalent to mx.split but allows for fractional segments.
python/mlx/nn/layers/distributed.py:30
↓ 2 callersMethod_split_dictionary
(self, gradients: dict)
python/mlx/optimizers/optimizers.py:184
↓ 2 callersMethod_validate_keys
(self, keys, strict)
python/mlx/nn/layers/base.py:456
↓ 2 callersMethodaccept
mlx/distributed/jaccl/lib/jaccl/tcp.cpp:126
↓ 2 callersFunctionaccept_connections
* Create a socket and accept one connection for each of the provided * addresses. */
mlx/distributed/ring/ring.cpp:331
↓ 2 callersMethodadd
mlx/distributed/ring/ring.cpp:279
↓ 2 callersMethodadd_graph_node
mlx/backend/cuda/device.cpp:424
↓ 2 callersFunctionaligned_alloc
mlx/3rdparty/pocketfft.h:156
↓ 2 callersFunctionaligned_dealloc
mlx/3rdparty/pocketfft.h:163
↓ 2 callersFunctionall
mlx/ops.cpp:1976
↓ 2 callersMethodall_gather
mlx/distributed/jaccl/lib/jaccl/ring.cpp:161
↓ 2 callersMethodallocate
mlx/3rdparty/pocketfft.h:608
↓ 2 callersFunctionallocate_workspace
mlx/backend/cuda/utils.cpp:82
↓ 2 callersFunctionappend_binary_kernels
mlx/backend/metal/jit_kernels.cpp:54
↓ 2 callersFunctionappend_indices_arg
mlx/backend/cuda/indexing.cpp:32
↓ 2 callersMethodappend_ptr
mlx/backend/cuda/jit_module.h:65
↓ 2 callersFunctionarccos
mlx/ops.cpp:3026
↓ 2 callersFunctionarccosh
mlx/ops.cpp:3073
↓ 2 callersFunctionarcsin
mlx/ops.cpp:3019
↓ 2 callersFunctionarcsinh
mlx/ops.cpp:3066
↓ 2 callersFunctionarctan2
mlx/ops.cpp:3040
↓ 2 callersFunctionarctanh
mlx/ops.cpp:3080
↓ 2 callersFunctionarray_from_list_impl
python/src/convert.cpp:364
↓ 2 callersMethodas_linear
Call the embedding layer as a linear layer. Use this for example when input embedding and output projection weights are tied
python/mlx/nn/layers/embedding.py:34
↓ 2 callersFunctionat_least_float
mlx/linalg.cpp:39
↓ 2 callersMethodatomic_update
mlx/backend/metal/kernels/reduction/ops.h:87
↓ 2 callersMethodbarrier
mlx/distributed/jaccl/lib/jaccl/rdma.h:336
↓ 2 callersFunctionbench
(f, a, b, b_prime)
benchmarks/python/conv3d_bench.py:13
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv2d_transpose_bench_cpu.py:14
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv_transpose_bench.py:16
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv_bench.py:19
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv1d_bench.py:19
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv3d_bench_cpu.py:15
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv_unaligned_bench.py:13
↓ 2 callersFunctionbench
(f, *args)
benchmarks/python/sdpa_bench.py:20
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv3d_transpose_bench_cpu.py:15
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/conv2d_bench_cpu.py:15
↓ 2 callersFunctionbench
(f, a, b)
benchmarks/python/blas/bench_gemm.py:21
↓ 2 callersFunctionbench_lens
(in_vec_len, out_vec_len, np_dtype, transpose=False)
benchmarks/python/blas/bench_gemv.py:86
↓ 2 callersFunctionbench_mlx
(fn, warmup, iters)
benchmarks/python/block_masked_mm_bench.py:54
↓ 2 callersFunctionbroadcast_vjp
mlx/primitives.cpp:812
↓ 2 callersMethodbuild
mlx/backend/cuda/compiled.cpp:25
↓ 2 callersFunctionbuild_sdpa_cache_key
mlx/backend/cuda/scaled_dot_product_attention.cpp:142
↓ 2 callersMethodcache_size
mlx/backend/common/buffer_cache.h:101
↓ 2 callersMethodcall_impl
python/src/transforms.cpp:449
↓ 2 callersMethodcapture_context
mlx/backend/cuda/device.h:42
↓ 2 callersFunctionceildiv
mlx/backend/cpu/gemms/simd_gemm.h:8
↓ 2 callersFunctioncheck_float_or_complex
mlx/linalg.cpp:30
↓ 2 callersFunctioncheck_kernel_threadgroup_size
mlx/backend/metal/utils.h:84
↓ 2 callersFunctioncheck_shape_dim
python/src/convert.cpp:17
↓ 2 callersFunctioncheck_valid_mesh
(hosts, connectivity, strict=True)
python/mlx/_distributed_utils/config.py:337
↓ 2 callersFunctionclear_cache
mlx/backend/cuda/allocator.cpp:442
↓ 2 callersFunctioncompile_clear_cache
mlx/compile.cpp:1196
↓ 2 callersFunctioncompile_dfs
Traverses the graph to build a tape and a map of array ids to their parents
mlx/compile.cpp:420
↓ 2 callersFunctioncompile_erase
mlx/compile.cpp:1192
↓ 2 callersFunctioncompile_replace
mlx/compile.cpp:1026
↓ 2 callersFunctioncompile_simplify
Simplify the tape. Note, this function modifies in-place both the tape, the parents map to remove orphaned arrays, and potentially the outputs
mlx/compile.cpp:569
↓ 2 callersFunctioncompile_trace
mlx/compile.cpp:398
↓ 2 callersFunctioncompiled_use_large_index
mlx/backend/common/compiled.cpp:224
↓ 2 callersFunctioncompute_bluestein_constants
Bluestein
mlx/backend/metal/fft.cpp:303
↓ 2 callersFunctioncompute_cost_and_scaling
mlx/einsum.cpp:141
↓ 2 callersFunctioncompute_dynamic_offset
mlx/backend/cpu/primitives.cpp:31
↓ 2 callersFunctioncompute_dynamic_offset
mlx/backend/cuda/slicing.cpp:44
↓ 2 callersMethodconcurrent_managed_access
mlx/backend/cuda/device.h:182
↓ 2 callersFunctionconfigure_jaccl
(args, hosts, ips, sshinfo)
python/mlx/_distributed_utils/config.py:465
↓ 2 callersFunctionconfigure_jaccl_ring
(args, hosts, ips, ring, sshinfo)
python/mlx/_distributed_utils/config.py:486
↓ 2 callersFunctionconfigure_ring
(args, hosts, ips, ring, sshinfo)
python/mlx/_distributed_utils/config.py:445
↓ 2 callersMethodcontiguous_suffix
mlx/backend/common/utils.h:145
↓ 2 callersFunctionconv_transpose2d
2D transposed convolution with a filter */
mlx/ops.cpp:4212
↓ 2 callersFunctionconv_transpose3d
3D transposed convolution with a filter */
mlx/ops.cpp:4233
↓ 2 callersFunctioncount_graph_nodes
tests/autograd_tests.cpp:22
↓ 2 callersMethodcpu_value
mlx/backend/metal/fence.cpp:41
↓ 2 callersFunctioncublas_bin_dir
mlx/backend/cuda/delayload.cpp:18
↓ 2 callersFunctioncustom_function
mlx/transforms.cpp:960
↓ 2 callersFunctioncustom_vjp
mlx/transforms.cpp:1057
↓ 2 callersFunctiondecompose_hadamard
mlx/backend/common/hadamard.h:85
↓ 2 callersFunctiondepends
mlx/ops.cpp:6069
↓ 2 callersMethoddetach
mlx/distributed/jaccl/lib/jaccl/tcp.cpp:84
↓ 2 callersFunctiondetach_event
mlx/array.h:428
↓ 2 callersFunctiondilate_size
Conv helpers
mlx/primitives.cpp:1308
↓ 2 callersFunctiondispatch_conv_2D_gpu
mlx/backend/metal/conv.cpp:970
↓ 2 callersFunctiondispatch_qmv
mlx/backend/metal/quantized.cpp:1365
↓ 2 callersFunctiondo_attention
(f, q, k, v, scale, mask=None, transpose=False)
python/tests/test_fast_sdpa.py:67
↓ 2 callersFunctioneig
mlx/linalg.cpp:559
↓ 2 callersFunctioneinsum_path_helper
mlx/einsum.cpp:630
↓ 2 callersMethodelems
mlx/backend/metal/kernels/steel/attn/nax.h:631
↓ 2 callersFunctionelu
r"""Applies the Exponential Linear Unit. Simply ``mx.where(x > 0, x, alpha * (mx.exp(x) - 1))``.
python/mlx/nn/layers/activations.py:75
↓ 2 callersMethodempty
mlx/3rdparty/pocketfft.h:597
↓ 2 callersMethodencode_capturing
mlx/backend/cuda/cudnn_utils.cpp:138
↓ 2 callersMethodencode_graph
mlx/backend/cuda/cudnn_utils.cpp:111
↓ 2 callersFunctionend
mlx/array.h:192
↓ 2 callersFunctionensure_batch_contiguous
mlx/backend/metal/matmul.cpp:54
↓ 2 callersFunctionensure_batch_contiguous
mlx/backend/cuda/matmul.cpp:35
← previousnext →901–1,000 of 4,946, ranked by callers