mlx/docs/src/python/ops.rst

.. _ops:

Operations
==========

.. currentmodule:: mlx.core

.. autosummary::
  :toctree: _autosummary

   abs
   add
   addmm
   all
   allclose
   any
   arange
   arccos
   arccosh
   arcsin
   arcsinh
   arctan
   arctan2
   arctanh
   argmax
   argmin
   argpartition
   argsort
   array_equal
   as_strided
   atleast_1d
   atleast_2d
   atleast_3d
   bitwise_and
   bitwise_or
   bitwise_xor
   block_masked_mm
   broadcast_to
   ceil
   clip
   concatenate
   conj
   conjugate
   convolve
   conv1d
   conv2d
   conv3d
   conv_transpose1d
   conv_transpose2d
   conv_transpose3d
   conv_general
   cos
   cosh
   cummax
   cummin
   cumprod
   cumsum
   degrees
   dequantize
   diag
   diagonal
   divide
   divmod
   einsum
   einsum_path
   equal
   erf
   erfinv
   exp
   expm1
   expand_dims
   eye
   flatten
   floor
   floor_divide
   full
   gather_mm
   gather_qmm
   greater
   greater_equal
   hadamard_transform
   identity
   inner
   isfinite
   isclose
   isinf
   isnan
   isneginf
   isposinf
   issubdtype
   left_shift
   less
   less_equal
   linspace
   load
   log
   log2
   log10
   log1p
   logaddexp
   logical_not
   logical_and
   logical_or
   logsumexp
   matmul
   max
   maximum
   mean
   meshgrid
   min
   minimum
   moveaxis
   multiply
   nan_to_num
   negative
   not_equal
   ones
   ones_like
   outer
   partition
   pad
   power
   prod
   put_along_axis
   quantize
   quantized_matmul
   radians
   reciprocal
   remainder
   repeat
   reshape
   right_shift
   roll
   round
   rsqrt
   save
   savez
   savez_compressed
   save_gguf
   save_safetensors
   sigmoid
   sign
   sin
   sinh
   softmax
   sort
   split
   sqrt
   square
   squeeze
   stack
   std
   stop_gradient
   subtract
   sum
   swapaxes
   take
   take_along_axis
   tan
   tanh
   tensordot
   tile
   topk
   trace
   transpose
   tri
   tril
   triu
   var
   view
   where
   zeros
   zeros_like
angelos's commit files 2023-11-29 10:42:59 -08:00			`.. _ops:`

			`Operations`
			`==========`

			`.. currentmodule:: mlx.core`

std and expm1 (#973) * std and expm1 * actually add expm1 * fix linux * fix vjp * relax tol for linux test * Add it to the compilable primitives --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> 2024-04-08 14:26:01 -07:00			`.. autosummary::`
angelos's commit files 2023-11-29 10:42:59 -08:00			`:toctree: _autosummary`

			`abs`
			`add`
Add missing docs for some ops (#1110) 2024-05-14 21:09:05 +08:00			`addmm`
angelos's commit files 2023-11-29 10:42:59 -08:00			`all`
std and expm1 (#973) * std and expm1 * actually add expm1 * fix linux * fix vjp * relax tol for linux test * Add it to the compilable primitives --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> 2024-04-08 14:26:01 -07:00			`allclose`
angelos's commit files 2023-11-29 10:42:59 -08:00			`any`
			`arange`
			`arccos`
			`arccosh`
			`arcsin`
			`arcsinh`
			`arctan`
Added ArcTan2 operation (#1079) * Added ArcTan2 operation * Cleanup, bug fixes from code review * Minor cleanup, fixed Linux tests 2024-05-08 11:35:15 -04:00			`arctan2`
angelos's commit files 2023-11-29 10:42:59 -08:00			`arctanh`
			`argmax`
			`argmin`
			`argpartition`
			`argsort`
			`array_equal`
Add missing docs for some ops (#1110) 2024-05-14 21:09:05 +08:00			`as_strided`
Added support for atleast_1d, atleast_2d, atleast_3d (#694) 2024-02-19 12:40:52 -05:00			`atleast_1d`
			`atleast_2d`
			`atleast_3d`
Add bitwise ops (#1037) * bitwise ops * fix tests 2024-04-26 22:03:42 -07:00			`bitwise_and`
			`bitwise_or`
			`bitwise_xor`
Masked mm (#978) * Add block masked matmul op and primitive 2024-04-16 14:45:39 -07:00			`block_masked_mm`
Add bitwise ops (#1037) * bitwise ops * fix tests 2024-04-26 22:03:42 -07:00			`broadcast_to`
Floor and Ceil (#150) * Implements Floor and Ceil Ops 2023-12-14 19:00:23 +01:00			`ceil`
Added clip function (#159) * Added clip * Added Python bindings * Formatting * Added cpp tests * Added Python tests * python bindings work * rebase --------- Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-17 22:00:29 -06:00			`clip`
angelos's commit files 2023-11-29 10:42:59 -08:00			`concatenate`
Add conjugate operator (#1100) * cpu and gpu impl * add mx.conj and array.conj() --------- Co-authored-by: Alex Barron <abarron22@apple.com> 2024-05-10 07:22:20 -07:00			`conj`
			`conjugate`
angelos's commit files 2023-11-29 10:42:59 -08:00			`convolve`
			`conv1d`
			`conv2d`
Bump extension MLX version (#1350) * Bump extension MLX version * fix some docs nits 2024-08-23 12:38:34 -07:00			`conv3d`
Transposed Convolution (#1245) * initial implementation for conv_transpose ran pre-commit implemented conv_transpose updated conv_general docstring updated conv_general docstring updated code comments removed commented run_conv_checks updated acknowledgments added missing entry to ops.rst added op to nn.layers resolved merge conflicts * removed ConvolutionTranspose primitive as suggested by reviewer removed ConvolutionTranspose primitive as suggested by reviewer * remove transpose flag, add another test --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-09-07 04:52:38 +02:00			`conv_transpose1d`
			`conv_transpose2d`
			`conv_transpose3d`
Convolution update (#651) * Init steel conv and update Conv primitive * Update slow CPU implementation to support flipping and input dilation winograd conv routing Co-authored-by: Awni Hannun <awni@apple.com> 2024-02-28 20:11:16 -08:00			`conv_general`
angelos's commit files 2023-11-29 10:42:59 -08:00			`cos`
			`cosh`
post nanobind docs fixes and some updates (#889) * post nanobind docs fixes and some updates * one more doc nit * fix for stubs and latex 2024-03-24 15:03:27 -07:00			`cummax`
			`cummin`
			`cumprod`
			`cumsum`
Adds radians and degrees (#1011) 2024-04-22 13:17:49 -05:00			`degrees`
Adds C++ and nn quantization utilities (#230) * Add C++ de-/quantize ops * Add quantize functions to the docs and tests * Add a QuantizedLinear module 2023-12-20 14:17:38 -08:00			`dequantize`
Implement diagonal operator (#562) * Implement diagonal operator This implements mx.diagonal in operator level, inspired by @ManishAradwad. * added `mx.diag` with tests * corrected few things * nits in bindings * updates to diag --------- Co-authored-by: ManishAradwad <manisharadwad@gmail.com> Co-authored-by: Awni Hannun <awni@apple.com> 2024-01-30 11:45:48 -06:00			`diag`
			`diagonal`
angelos's commit files 2023-11-29 10:42:59 -08:00			`divide`
Multi output primitives (#330) * Multi-output primitives --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> 2024-01-08 16:39:08 -08:00			`divmod`
Einsum (#1269) * einsum initial * fix comma break * sum axis was wrong * small cleanups * python binding * changed bindings to resemble numpy * remove todo comment * comment changes * add count of operands/inputs * fail fast if operands list is empty * ignore comma if no output * einsum path matching numpy * getting somewhere with path * remove print * it passes the first test * moved einsum tests to seperate file * seperated einsum path * moved einsum naive * remove space from equation * fast fail if no operands passed * update tests and remove printf * small cleanup * some more cleanups * removed python helper file * ack * utilize std for finding min in vector * duplicate def * remove the tuple as it was unreadable * moved einsum_naive back to ops * remaining isn't needed * avoid creating another set * cleanup * greedy path, start of naive einsum * more einsum * fix some bugs * some more fixes, tests pass * benchmark * some simplify * fix einsum and test Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> * add a bunch more tests and fix a bunch more bugs * some docs nits --------- Co-authored-by: dc-dc-dc <dgcruz983@gmail.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> 2024-07-25 09:36:44 -07:00			`einsum`
			`einsum_path`
angelos's commit files 2023-11-29 10:42:59 -08:00			`equal`
			`erf`
			`erfinv`
			`exp`
std and expm1 (#973) * std and expm1 * actually add expm1 * fix linux * fix vjp * relax tol for linux test * Add it to the compilable primitives --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> 2024-04-08 14:26:01 -07:00			`expm1`
angelos's commit files 2023-11-29 10:42:59 -08:00			`expand_dims`
Added eye/identity ops (#119) `eye` and `identity` C++ and Python ops 2023-12-11 12:38:17 -08:00			`eye`
implemented Flatten Module (#149) * implemented flatten op --------- Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-17 06:54:37 +01:00			`flatten`
Implements divide for integer types and adds floor_divide op (#228) * Add floor_divide * Add floor_divide to the tests * Add floor_divide to the docs 2023-12-19 20:12:19 -08:00			`floor`
			`floor_divide`
angelos's commit files 2023-11-29 10:42:59 -08:00			`full`
Rename block sparse (#1149) * block_sparse_mm to gather_mm * rename * nit * nit 2024-05-22 07:48:34 -07:00			`gather_mm`
			`gather_qmm`
angelos's commit files 2023-11-29 10:42:59 -08:00			`greater`
			`greater_equal`
Fast Hadamard Transform (#1249) * Working hadamard for powers of 2 * working for m2^k add scale and check contiguity * add size check * clean up * fix test * add grads + vmap * gpu only * skip on linux * test typo * add cpu impl * remove gpu only tests * fix linux build + add is_equivalent 2024-07-09 20:39:01 -07:00			`hadamard_transform`
Added eye/identity ops (#119) `eye` and `identity` C++ and Python ops 2023-12-11 12:38:17 -08:00			`identity`
Add inner / outer op (#348) * inner / outer impl * python tests * ops list and ack * updated descriptions * use test helper * removed dtype check and flatten outer to 1-D * updated docs * just use the reshape to flatten 2024-01-07 12:01:09 -05:00			`inner`
add std as method (#1387) * add std as method * add std as method 2024-09-01 19:49:16 -07:00			`isfinite`
bindings for memory info (#761) * bindings for memory info * update api * keep cache low if requested * fix default * nit in ops error 2024-03-01 19:51:58 -08:00			`isclose`
add numeric type hierarchy and issubdtype as well as a set_dtype meth… (#427) * add numeric type hierarchy and issubdtype as well as a set_dtype method to nn.Module with predicate numeric type hierarchy and issubtype is compatible to the [numpy hierarchy](https://github.com/numpy/numpy/blob/220f0ab2c5958f9d49b9a28c69a760657b06d9b9/numpy/_core/numerictypes.py#L42). Closes #285. * nits in docs * unify type category checking * nits in docs * nits in docs * more docs nits * fix callable type --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-03-25 20:32:59 +01:00			`isinf`
Add isnan (#423) 2024-01-12 14:16:48 -05:00			`isnan`
implemented isposinf and isneginf in one PR (#470) * ran precommit * updated docs 2024-01-16 20:18:07 +05:30			`isneginf`
add numeric type hierarchy and issubdtype as well as a set_dtype meth… (#427) * add numeric type hierarchy and issubdtype as well as a set_dtype method to nn.Module with predicate numeric type hierarchy and issubtype is compatible to the [numpy hierarchy](https://github.com/numpy/numpy/blob/220f0ab2c5958f9d49b9a28c69a760657b06d9b9/numpy/_core/numerictypes.py#L42). Closes #285. * nits in docs * unify type category checking * nits in docs * nits in docs * more docs nits * fix callable type --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-03-25 20:32:59 +01:00			`isposinf`
Add missing docs for some ops (#1110) 2024-05-14 21:09:05 +08:00			`issubdtype`
Add bitwise ops (#1037) * bitwise ops * fix tests 2024-04-26 22:03:42 -07:00			`left_shift`
angelos's commit files 2023-11-29 10:42:59 -08:00			`less`
			`less_equal`
Added linspace (#181) * linspace ops support --------- Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-18 19:57:55 -08:00			`linspace`
angelos's commit files 2023-11-29 10:42:59 -08:00			`load`
			`log`
			`log2`
			`log10`
			`log1p`
			`logaddexp`
			`logical_not`
feat: add logicalAnd and logicalOR (#386) * feat: add logicalAnd and logicalOR * run pre-commit * Refactor logical_and and logical_or functions * Add acknowledgement * Add logical AND and logical OR operators * Refactor logical_and and logical_or functions * Add support for logical operators on bool arrays * Update mlx/ops.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Update mlx/ops.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Add logical AND and OR operators for arrays and scalars * Refactor vjp and jvp methods in primitives.cpp * Add overloaded operators for logical AND and OR * format --------- Co-authored-by: Awni Hannun <awni.hannun@gmail.com> Co-authored-by: Awni Hannun <awni@apple.com> 2024-01-08 19:00:05 +04:00			`logical_and`
			`logical_or`
angelos's commit files 2023-11-29 10:42:59 -08:00			`logsumexp`
			`matmul`
			`max`
			`maximum`
			`mean`
Add mx.meshgrid (#961) 2024-04-09 14:43:08 -04:00			`meshgrid`
angelos's commit files 2023-11-29 10:42:59 -08:00			`min`
			`minimum`
Add move and swap axis, and vmap for slice, concat, and gather (#158) * add move and swap axis, and vmap for slice, concat, and gather 2023-12-14 12:59:12 -08:00			`moveaxis`
angelos's commit files 2023-11-29 10:42:59 -08:00			`multiply`
[Issue #1187] Add nan_to_num function initial attempt (#1247) * initial attempt, working with wrong types * not compiling; mx.float16 and mx.bfloat16 tests added * fix nan to num * nit --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-07-25 17:57:37 +01:00			`nan_to_num`
angelos's commit files 2023-11-29 10:42:59 -08:00			`negative`
Expose function to clear memory cache (#1032) * expose function to clear memory cache * fix linux build * fix metal tests 2024-04-24 16:48:51 -07:00			`not_equal`
angelos's commit files 2023-11-29 10:42:59 -08:00			`ones`
			`ones_like`
Add inner / outer op (#348) * inner / outer impl * python tests * ops list and ack * updated descriptions * use test helper * removed dtype check and flatten outer to 1-D * updated docs * just use the reshape to flatten 2024-01-07 12:01:09 -05:00			`outer`
angelos's commit files 2023-11-29 10:42:59 -08:00			`partition`
			`pad`
Add missing docs for some ops (#1110) 2024-05-14 21:09:05 +08:00			`power`
angelos's commit files 2023-11-29 10:42:59 -08:00			`prod`
Put along axis + fixe for partition grad (#1430) * put along axis, fixes for partition grad * zeros for arg reduce 2024-09-23 10:03:38 -07:00			`put_along_axis`
Adds C++ and nn quantization utilities (#230) * Add C++ de-/quantize ops * Add quantize functions to the docs and tests * Add a QuantizedLinear module 2023-12-20 14:17:38 -08:00			`quantize`
			`quantized_matmul`
Adds radians and degrees (#1011) 2024-04-22 13:17:49 -05:00			`radians`
angelos's commit files 2023-11-29 10:42:59 -08:00			`reciprocal`
Add missing docs for some ops (#1110) 2024-05-14 21:09:05 +08:00			`remainder`
Add support for repeat (#278) * add repeat function * fix styling * optimizing repeat * fixed minor issues * not sure why that folder is there xD * fixed now for sure * test repeat not repeat test * Fixed --------- Co-authored-by: Bahaa Eddin tabbakha <bahaa@Bahaas-MacBook-Pro.local> 2023-12-28 00:11:38 +03:00			`repeat`
angelos's commit files 2023-11-29 10:42:59 -08:00			`reshape`
Add bitwise ops (#1037) * bitwise ops * fix tests 2024-04-26 22:03:42 -07:00			`right_shift`
Add the roll op (#1455) 2024-10-07 17:21:42 -07:00			`roll`
Adds round op and primitive (#203) 2023-12-18 11:32:48 -08:00			`round`
angelos's commit files 2023-11-29 10:42:59 -08:00			`rsqrt`
			`save`
			`savez`
			`savez_compressed`
GGUF support (#350) * Initial GGUF support for tensor fields. --------- Co-authored-by: Awni Hannun <awni@apple.com> 2024-01-10 16:22:48 -05:00			`save_gguf`
Safetensor support (#215) Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-27 05:06:55 -05:00			`save_safetensors`
angelos's commit files 2023-11-29 10:42:59 -08:00			`sigmoid`
			`sign`
			`sin`
			`sinh`
			`softmax`
			`sort`
			`split`
			`sqrt`
			`square`
			`squeeze`
Added mx.stack c++ frontend impl (#123) * stack C++ operation + python bindings 2023-12-14 16:21:19 -05:00			`stack`
std and expm1 (#973) * std and expm1 * actually add expm1 * fix linux * fix vjp * relax tol for linux test * Add it to the compilable primitives --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> 2024-04-08 14:26:01 -07:00			`std`
angelos's commit files 2023-11-29 10:42:59 -08:00			`stop_gradient`
			`subtract`
			`sum`
Add move and swap axis, and vmap for slice, concat, and gather (#158) * add move and swap axis, and vmap for slice, concat, and gather 2023-12-14 12:59:12 -08:00			`swapaxes`
angelos's commit files 2023-11-29 10:42:59 -08:00			`take`
			`take_along_axis`
			`tan`
			`tanh`
Add Tensordot op (#344) 2024-01-02 20:15:00 -05:00			`tensordot`
bindings for memory info (#761) * bindings for memory info * update api * keep cache low if requested * fix default * nit in ops error 2024-03-01 19:51:58 -08:00			`tile`
Pickle + dtype fix for numpy conversion (#763) * pickle + dtype fix for numpy conversion * fix getattribute on Module base * remove unused function * fix tests * add topk to ops * fix doc 2024-03-02 06:09:29 -08:00			`topk`
add mx.trace (#1143) (#1147) * working c++ trace implementation * updated throw + added overloads * added python binding for trace function * pre-commit reformatting * add trace to docs * resolve comments * remove to_stream call 2024-05-22 18:50:27 -04:00			`trace`
angelos's commit files 2023-11-29 10:42:59 -08:00			`transpose`
added tri / tril / triu (#170) * added tri / tril / triu * fixed tests * ctest tests * tri overload and simplified tests * changes from comment * more tests for m * ensure assert if not 2-D * remove broadcast_to * minor tweaks --------- Co-authored-by: Awni Hannun <awni@apple.com> 2023-12-15 20:30:34 -05:00			`tri`
			`tril`
			`triu`
angelos's commit files 2023-11-29 10:42:59 -08:00			`var`
Add view op (#1179) * add view primitive * nit * fix view 2024-06-04 08:05:27 -07:00			`view`
angelos's commit files 2023-11-29 10:42:59 -08:00			`where`
			`zeros`
			`zeros_like`