Search.setIndex({"docnames": ["cpp/ops", "dev/extensions", "dev/metal_debugger", "examples/linear_regression", "examples/llama-inference", "examples/mlp", "index", "install", "python/_autosummary/mlx.core.Device", "python/_autosummary/mlx.core.Dtype", "python/_autosummary/mlx.core.DtypeCategory", "python/_autosummary/mlx.core.abs", "python/_autosummary/mlx.core.add", "python/_autosummary/mlx.core.all", "python/_autosummary/mlx.core.allclose", "python/_autosummary/mlx.core.any", "python/_autosummary/mlx.core.arange", "python/_autosummary/mlx.core.arccos", "python/_autosummary/mlx.core.arccosh", "python/_autosummary/mlx.core.arcsin", "python/_autosummary/mlx.core.arcsinh", "python/_autosummary/mlx.core.arctan", "python/_autosummary/mlx.core.arctanh", "python/_autosummary/mlx.core.argmax", "python/_autosummary/mlx.core.argmin", "python/_autosummary/mlx.core.argpartition", "python/_autosummary/mlx.core.argsort", "python/_autosummary/mlx.core.array", "python/_autosummary/mlx.core.array.T", "python/_autosummary/mlx.core.array.abs", "python/_autosummary/mlx.core.array.all", "python/_autosummary/mlx.core.array.any", "python/_autosummary/mlx.core.array.argmax", "python/_autosummary/mlx.core.array.argmin", "python/_autosummary/mlx.core.array.astype", "python/_autosummary/mlx.core.array.at", "python/_autosummary/mlx.core.array.cos", "python/_autosummary/mlx.core.array.cummax", "python/_autosummary/mlx.core.array.cummin", "python/_autosummary/mlx.core.array.cumprod", "python/_autosummary/mlx.core.array.cumsum", "python/_autosummary/mlx.core.array.diag", "python/_autosummary/mlx.core.array.diagonal", "python/_autosummary/mlx.core.array.dtype", "python/_autosummary/mlx.core.array.exp", "python/_autosummary/mlx.core.array.flatten", "python/_autosummary/mlx.core.array.item", "python/_autosummary/mlx.core.array.itemsize", "python/_autosummary/mlx.core.array.log", "python/_autosummary/mlx.core.array.log10", "python/_autosummary/mlx.core.array.log1p", "python/_autosummary/mlx.core.array.log2", "python/_autosummary/mlx.core.array.logsumexp", "python/_autosummary/mlx.core.array.max", "python/_autosummary/mlx.core.array.mean", "python/_autosummary/mlx.core.array.min", "python/_autosummary/mlx.core.array.moveaxis", "python/_autosummary/mlx.core.array.nbytes", "python/_autosummary/mlx.core.array.ndim", "python/_autosummary/mlx.core.array.prod", "python/_autosummary/mlx.core.array.reciprocal", "python/_autosummary/mlx.core.array.reshape", "python/_autosummary/mlx.core.array.round", "python/_autosummary/mlx.core.array.rsqrt", "python/_autosummary/mlx.core.array.shape", "python/_autosummary/mlx.core.array.sin", "python/_autosummary/mlx.core.array.size", "python/_autosummary/mlx.core.array.split", "python/_autosummary/mlx.core.array.sqrt", "python/_autosummary/mlx.core.array.square", "python/_autosummary/mlx.core.array.squeeze", "python/_autosummary/mlx.core.array.sum", "python/_autosummary/mlx.core.array.swapaxes", "python/_autosummary/mlx.core.array.tolist", "python/_autosummary/mlx.core.array.transpose", "python/_autosummary/mlx.core.array.var", "python/_autosummary/mlx.core.array_equal", "python/_autosummary/mlx.core.atleast_1d", "python/_autosummary/mlx.core.atleast_2d", "python/_autosummary/mlx.core.atleast_3d", "python/_autosummary/mlx.core.broadcast_to", "python/_autosummary/mlx.core.ceil", "python/_autosummary/mlx.core.clip", "python/_autosummary/mlx.core.compile", "python/_autosummary/mlx.core.concatenate", "python/_autosummary/mlx.core.conv1d", "python/_autosummary/mlx.core.conv2d", "python/_autosummary/mlx.core.conv_general", "python/_autosummary/mlx.core.convolve", "python/_autosummary/mlx.core.cos", "python/_autosummary/mlx.core.cosh", "python/_autosummary/mlx.core.cummax", "python/_autosummary/mlx.core.cummin", "python/_autosummary/mlx.core.cumprod", "python/_autosummary/mlx.core.cumsum", "python/_autosummary/mlx.core.default_device", "python/_autosummary/mlx.core.default_stream", "python/_autosummary/mlx.core.dequantize", "python/_autosummary/mlx.core.diag", "python/_autosummary/mlx.core.diagonal", "python/_autosummary/mlx.core.disable_compile", "python/_autosummary/mlx.core.divide", "python/_autosummary/mlx.core.divmod", "python/_autosummary/mlx.core.enable_compile", "python/_autosummary/mlx.core.equal", "python/_autosummary/mlx.core.erf", "python/_autosummary/mlx.core.erfinv", "python/_autosummary/mlx.core.eval", "python/_autosummary/mlx.core.exp", "python/_autosummary/mlx.core.expand_dims", "python/_autosummary/mlx.core.expm1", "python/_autosummary/mlx.core.eye", "python/_autosummary/mlx.core.fast.layer_norm", "python/_autosummary/mlx.core.fast.rms_norm", "python/_autosummary/mlx.core.fast.rope", "python/_autosummary/mlx.core.fast.scaled_dot_product_attention", "python/_autosummary/mlx.core.fft.fft", "python/_autosummary/mlx.core.fft.fft2", "python/_autosummary/mlx.core.fft.fftn", "python/_autosummary/mlx.core.fft.ifft", "python/_autosummary/mlx.core.fft.ifft2", "python/_autosummary/mlx.core.fft.ifftn", "python/_autosummary/mlx.core.fft.irfft", "python/_autosummary/mlx.core.fft.irfft2", "python/_autosummary/mlx.core.fft.irfftn", "python/_autosummary/mlx.core.fft.rfft", "python/_autosummary/mlx.core.fft.rfft2", "python/_autosummary/mlx.core.fft.rfftn", "python/_autosummary/mlx.core.flatten", "python/_autosummary/mlx.core.floor", "python/_autosummary/mlx.core.floor_divide", "python/_autosummary/mlx.core.full", "python/_autosummary/mlx.core.grad", "python/_autosummary/mlx.core.greater", "python/_autosummary/mlx.core.greater_equal", "python/_autosummary/mlx.core.identity", "python/_autosummary/mlx.core.inner", "python/_autosummary/mlx.core.isclose", "python/_autosummary/mlx.core.isinf", "python/_autosummary/mlx.core.isnan", "python/_autosummary/mlx.core.isneginf", "python/_autosummary/mlx.core.isposinf", "python/_autosummary/mlx.core.issubdtype", "python/_autosummary/mlx.core.jvp", "python/_autosummary/mlx.core.less", "python/_autosummary/mlx.core.less_equal", "python/_autosummary/mlx.core.linalg.norm", "python/_autosummary/mlx.core.linalg.qr", "python/_autosummary/mlx.core.linspace", "python/_autosummary/mlx.core.load", "python/_autosummary/mlx.core.log", "python/_autosummary/mlx.core.log10", "python/_autosummary/mlx.core.log1p", "python/_autosummary/mlx.core.log2", "python/_autosummary/mlx.core.logaddexp", "python/_autosummary/mlx.core.logical_and", "python/_autosummary/mlx.core.logical_not", "python/_autosummary/mlx.core.logical_or", "python/_autosummary/mlx.core.logsumexp", "python/_autosummary/mlx.core.matmul", "python/_autosummary/mlx.core.max", "python/_autosummary/mlx.core.maximum", "python/_autosummary/mlx.core.mean", "python/_autosummary/mlx.core.meshgrid", "python/_autosummary/mlx.core.metal.get_active_memory", "python/_autosummary/mlx.core.metal.get_cache_memory", "python/_autosummary/mlx.core.metal.get_peak_memory", "python/_autosummary/mlx.core.metal.is_available", "python/_autosummary/mlx.core.metal.set_cache_limit", "python/_autosummary/mlx.core.metal.set_memory_limit", "python/_autosummary/mlx.core.metal.start_capture", "python/_autosummary/mlx.core.metal.stop_capture", "python/_autosummary/mlx.core.min", "python/_autosummary/mlx.core.minimum", "python/_autosummary/mlx.core.moveaxis", "python/_autosummary/mlx.core.multiply", "python/_autosummary/mlx.core.negative", "python/_autosummary/mlx.core.new_stream", "python/_autosummary/mlx.core.ones", "python/_autosummary/mlx.core.ones_like", "python/_autosummary/mlx.core.outer", "python/_autosummary/mlx.core.pad", "python/_autosummary/mlx.core.partition", "python/_autosummary/mlx.core.prod", "python/_autosummary/mlx.core.quantize", "python/_autosummary/mlx.core.quantized_matmul", "python/_autosummary/mlx.core.random.bernoulli", "python/_autosummary/mlx.core.random.categorical", "python/_autosummary/mlx.core.random.gumbel", "python/_autosummary/mlx.core.random.key", "python/_autosummary/mlx.core.random.multivariate_normal", "python/_autosummary/mlx.core.random.normal", "python/_autosummary/mlx.core.random.randint", "python/_autosummary/mlx.core.random.seed", "python/_autosummary/mlx.core.random.split", "python/_autosummary/mlx.core.random.truncated_normal", "python/_autosummary/mlx.core.random.uniform", "python/_autosummary/mlx.core.reciprocal", "python/_autosummary/mlx.core.repeat", "python/_autosummary/mlx.core.reshape", "python/_autosummary/mlx.core.round", "python/_autosummary/mlx.core.rsqrt", "python/_autosummary/mlx.core.save", "python/_autosummary/mlx.core.save_gguf", "python/_autosummary/mlx.core.save_safetensors", "python/_autosummary/mlx.core.savez", "python/_autosummary/mlx.core.savez_compressed", "python/_autosummary/mlx.core.set_default_device", "python/_autosummary/mlx.core.set_default_stream", "python/_autosummary/mlx.core.sigmoid", "python/_autosummary/mlx.core.sign", "python/_autosummary/mlx.core.sin", "python/_autosummary/mlx.core.sinh", "python/_autosummary/mlx.core.softmax", "python/_autosummary/mlx.core.sort", "python/_autosummary/mlx.core.split", "python/_autosummary/mlx.core.sqrt", "python/_autosummary/mlx.core.square", "python/_autosummary/mlx.core.squeeze", "python/_autosummary/mlx.core.stack", "python/_autosummary/mlx.core.std", "python/_autosummary/mlx.core.stop_gradient", "python/_autosummary/mlx.core.stream", "python/_autosummary/mlx.core.subtract", "python/_autosummary/mlx.core.sum", "python/_autosummary/mlx.core.swapaxes", "python/_autosummary/mlx.core.take", "python/_autosummary/mlx.core.take_along_axis", "python/_autosummary/mlx.core.tan", "python/_autosummary/mlx.core.tanh", "python/_autosummary/mlx.core.tensordot", "python/_autosummary/mlx.core.tile", "python/_autosummary/mlx.core.topk", "python/_autosummary/mlx.core.transpose", "python/_autosummary/mlx.core.tri", "python/_autosummary/mlx.core.tril", "python/_autosummary/mlx.core.triu", "python/_autosummary/mlx.core.value_and_grad", "python/_autosummary/mlx.core.var", "python/_autosummary/mlx.core.vjp", "python/_autosummary/mlx.core.vmap", "python/_autosummary/mlx.core.where", "python/_autosummary/mlx.core.zeros", "python/_autosummary/mlx.core.zeros_like", "python/_autosummary/mlx.nn.value_and_grad", "python/_autosummary/mlx.utils.tree_flatten", "python/_autosummary/mlx.utils.tree_map", "python/_autosummary/mlx.utils.tree_unflatten", "python/_autosummary/stream_class", "python/array", "python/data_types", "python/devices_and_streams", "python/fast", "python/fft", "python/linalg", "python/metal", "python/nn", "python/nn/_autosummary/mlx.nn.ALiBi", "python/nn/_autosummary/mlx.nn.AvgPool1d", "python/nn/_autosummary/mlx.nn.AvgPool2d", "python/nn/_autosummary/mlx.nn.BatchNorm", "python/nn/_autosummary/mlx.nn.Conv1d", "python/nn/_autosummary/mlx.nn.Conv2d", "python/nn/_autosummary/mlx.nn.Dropout", "python/nn/_autosummary/mlx.nn.Dropout2d", "python/nn/_autosummary/mlx.nn.Dropout3d", "python/nn/_autosummary/mlx.nn.Embedding", "python/nn/_autosummary/mlx.nn.GELU", "python/nn/_autosummary/mlx.nn.GRU", "python/nn/_autosummary/mlx.nn.GroupNorm", "python/nn/_autosummary/mlx.nn.InstanceNorm", "python/nn/_autosummary/mlx.nn.LSTM", "python/nn/_autosummary/mlx.nn.LayerNorm", "python/nn/_autosummary/mlx.nn.Linear", "python/nn/_autosummary/mlx.nn.MaxPool1d", "python/nn/_autosummary/mlx.nn.MaxPool2d", "python/nn/_autosummary/mlx.nn.Mish", "python/nn/_autosummary/mlx.nn.Module.apply", "python/nn/_autosummary/mlx.nn.Module.apply_to_modules", "python/nn/_autosummary/mlx.nn.Module.children", "python/nn/_autosummary/mlx.nn.Module.eval", "python/nn/_autosummary/mlx.nn.Module.filter_and_map", "python/nn/_autosummary/mlx.nn.Module.freeze", "python/nn/_autosummary/mlx.nn.Module.leaf_modules", "python/nn/_autosummary/mlx.nn.Module.load_weights", "python/nn/_autosummary/mlx.nn.Module.modules", "python/nn/_autosummary/mlx.nn.Module.named_modules", "python/nn/_autosummary/mlx.nn.Module.parameters", "python/nn/_autosummary/mlx.nn.Module.save_weights", "python/nn/_autosummary/mlx.nn.Module.set_dtype", "python/nn/_autosummary/mlx.nn.Module.state", "python/nn/_autosummary/mlx.nn.Module.train", "python/nn/_autosummary/mlx.nn.Module.trainable_parameters", "python/nn/_autosummary/mlx.nn.Module.training", "python/nn/_autosummary/mlx.nn.Module.unfreeze", "python/nn/_autosummary/mlx.nn.Module.update", "python/nn/_autosummary/mlx.nn.Module.update_modules", "python/nn/_autosummary/mlx.nn.MultiHeadAttention", "python/nn/_autosummary/mlx.nn.PReLU", "python/nn/_autosummary/mlx.nn.QuantizedLinear", "python/nn/_autosummary/mlx.nn.RMSNorm", "python/nn/_autosummary/mlx.nn.RNN", "python/nn/_autosummary/mlx.nn.ReLU", "python/nn/_autosummary/mlx.nn.RoPE", "python/nn/_autosummary/mlx.nn.SELU", "python/nn/_autosummary/mlx.nn.Sequential", "python/nn/_autosummary/mlx.nn.SiLU", "python/nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding", "python/nn/_autosummary/mlx.nn.Softshrink", "python/nn/_autosummary/mlx.nn.Step", "python/nn/_autosummary/mlx.nn.Transformer", "python/nn/_autosummary/mlx.nn.Upsample", "python/nn/_autosummary/mlx.nn.init.constant", "python/nn/_autosummary/mlx.nn.init.glorot_normal", "python/nn/_autosummary/mlx.nn.init.glorot_uniform", "python/nn/_autosummary/mlx.nn.init.he_normal", "python/nn/_autosummary/mlx.nn.init.he_uniform", "python/nn/_autosummary/mlx.nn.init.identity", "python/nn/_autosummary/mlx.nn.init.normal", "python/nn/_autosummary/mlx.nn.init.uniform", "python/nn/_autosummary_functions/mlx.nn.elu", "python/nn/_autosummary_functions/mlx.nn.gelu", "python/nn/_autosummary_functions/mlx.nn.gelu_approx", "python/nn/_autosummary_functions/mlx.nn.gelu_fast_approx", "python/nn/_autosummary_functions/mlx.nn.glu", "python/nn/_autosummary_functions/mlx.nn.hardswish", "python/nn/_autosummary_functions/mlx.nn.leaky_relu", "python/nn/_autosummary_functions/mlx.nn.log_sigmoid", "python/nn/_autosummary_functions/mlx.nn.log_softmax", "python/nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy", "python/nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss", "python/nn/_autosummary_functions/mlx.nn.losses.cross_entropy", "python/nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss", "python/nn/_autosummary_functions/mlx.nn.losses.hinge_loss", "python/nn/_autosummary_functions/mlx.nn.losses.huber_loss", "python/nn/_autosummary_functions/mlx.nn.losses.kl_div_loss", "python/nn/_autosummary_functions/mlx.nn.losses.l1_loss", "python/nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss", "python/nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss", "python/nn/_autosummary_functions/mlx.nn.losses.mse_loss", "python/nn/_autosummary_functions/mlx.nn.losses.nll_loss", "python/nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss", "python/nn/_autosummary_functions/mlx.nn.losses.triplet_loss", "python/nn/_autosummary_functions/mlx.nn.mish", "python/nn/_autosummary_functions/mlx.nn.prelu", "python/nn/_autosummary_functions/mlx.nn.relu", "python/nn/_autosummary_functions/mlx.nn.relu6", "python/nn/_autosummary_functions/mlx.nn.selu", "python/nn/_autosummary_functions/mlx.nn.sigmoid", "python/nn/_autosummary_functions/mlx.nn.silu", "python/nn/_autosummary_functions/mlx.nn.softmax", "python/nn/_autosummary_functions/mlx.nn.softplus", "python/nn/_autosummary_functions/mlx.nn.softshrink", "python/nn/_autosummary_functions/mlx.nn.step", "python/nn/_autosummary_functions/mlx.nn.tanh", "python/nn/functions", "python/nn/init", "python/nn/layers", "python/nn/losses", "python/nn/module", "python/ops", "python/optimizers", "python/optimizers/_autosummary/mlx.optimizers.AdaDelta", "python/optimizers/_autosummary/mlx.optimizers.Adafactor", "python/optimizers/_autosummary/mlx.optimizers.Adagrad", "python/optimizers/_autosummary/mlx.optimizers.Adam", "python/optimizers/_autosummary/mlx.optimizers.AdamW", "python/optimizers/_autosummary/mlx.optimizers.Adamax", "python/optimizers/_autosummary/mlx.optimizers.Lion", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.apply_gradients", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.init", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.state", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.update", "python/optimizers/_autosummary/mlx.optimizers.RMSprop", "python/optimizers/_autosummary/mlx.optimizers.SGD", "python/optimizers/_autosummary/mlx.optimizers.cosine_decay", "python/optimizers/_autosummary/mlx.optimizers.exponential_decay", "python/optimizers/_autosummary/mlx.optimizers.join_schedules", "python/optimizers/_autosummary/mlx.optimizers.linear_schedule", "python/optimizers/_autosummary/mlx.optimizers.step_decay", "python/optimizers/common_optimizers", "python/optimizers/optimizer", "python/optimizers/schedulers", "python/random", "python/transforms", "python/tree_utils", "usage/compile", "usage/function_transforms", "usage/indexing", "usage/lazy_evaluation", "usage/numpy", "usage/quick_start", "usage/saving_and_loading", "usage/unified_memory", "usage/using_streams"], "filenames": ["cpp/ops.rst", "dev/extensions.rst", "dev/metal_debugger.rst", "examples/linear_regression.rst", "examples/llama-inference.rst", "examples/mlp.rst", "index.rst", "install.rst", "python/_autosummary/mlx.core.Device.rst", "python/_autosummary/mlx.core.Dtype.rst", "python/_autosummary/mlx.core.DtypeCategory.rst", "python/_autosummary/mlx.core.abs.rst", "python/_autosummary/mlx.core.add.rst", "python/_autosummary/mlx.core.all.rst", "python/_autosummary/mlx.core.allclose.rst", "python/_autosummary/mlx.core.any.rst", "python/_autosummary/mlx.core.arange.rst", "python/_autosummary/mlx.core.arccos.rst", "python/_autosummary/mlx.core.arccosh.rst", "python/_autosummary/mlx.core.arcsin.rst", "python/_autosummary/mlx.core.arcsinh.rst", "python/_autosummary/mlx.core.arctan.rst", "python/_autosummary/mlx.core.arctanh.rst", "python/_autosummary/mlx.core.argmax.rst", "python/_autosummary/mlx.core.argmin.rst", "python/_autosummary/mlx.core.argpartition.rst", "python/_autosummary/mlx.core.argsort.rst", "python/_autosummary/mlx.core.array.rst", "python/_autosummary/mlx.core.array.T.rst", "python/_autosummary/mlx.core.array.abs.rst", "python/_autosummary/mlx.core.array.all.rst", "python/_autosummary/mlx.core.array.any.rst", "python/_autosummary/mlx.core.array.argmax.rst", "python/_autosummary/mlx.core.array.argmin.rst", "python/_autosummary/mlx.core.array.astype.rst", "python/_autosummary/mlx.core.array.at.rst", "python/_autosummary/mlx.core.array.cos.rst", "python/_autosummary/mlx.core.array.cummax.rst", "python/_autosummary/mlx.core.array.cummin.rst", "python/_autosummary/mlx.core.array.cumprod.rst", "python/_autosummary/mlx.core.array.cumsum.rst", "python/_autosummary/mlx.core.array.diag.rst", "python/_autosummary/mlx.core.array.diagonal.rst", "python/_autosummary/mlx.core.array.dtype.rst", "python/_autosummary/mlx.core.array.exp.rst", "python/_autosummary/mlx.core.array.flatten.rst", "python/_autosummary/mlx.core.array.item.rst", "python/_autosummary/mlx.core.array.itemsize.rst", "python/_autosummary/mlx.core.array.log.rst", "python/_autosummary/mlx.core.array.log10.rst", "python/_autosummary/mlx.core.array.log1p.rst", "python/_autosummary/mlx.core.array.log2.rst", "python/_autosummary/mlx.core.array.logsumexp.rst", "python/_autosummary/mlx.core.array.max.rst", "python/_autosummary/mlx.core.array.mean.rst", "python/_autosummary/mlx.core.array.min.rst", "python/_autosummary/mlx.core.array.moveaxis.rst", "python/_autosummary/mlx.core.array.nbytes.rst", "python/_autosummary/mlx.core.array.ndim.rst", "python/_autosummary/mlx.core.array.prod.rst", "python/_autosummary/mlx.core.array.reciprocal.rst", "python/_autosummary/mlx.core.array.reshape.rst", "python/_autosummary/mlx.core.array.round.rst", "python/_autosummary/mlx.core.array.rsqrt.rst", "python/_autosummary/mlx.core.array.shape.rst", "python/_autosummary/mlx.core.array.sin.rst", "python/_autosummary/mlx.core.array.size.rst", "python/_autosummary/mlx.core.array.split.rst", "python/_autosummary/mlx.core.array.sqrt.rst", "python/_autosummary/mlx.core.array.square.rst", "python/_autosummary/mlx.core.array.squeeze.rst", "python/_autosummary/mlx.core.array.sum.rst", "python/_autosummary/mlx.core.array.swapaxes.rst", "python/_autosummary/mlx.core.array.tolist.rst", "python/_autosummary/mlx.core.array.transpose.rst", "python/_autosummary/mlx.core.array.var.rst", "python/_autosummary/mlx.core.array_equal.rst", "python/_autosummary/mlx.core.atleast_1d.rst", "python/_autosummary/mlx.core.atleast_2d.rst", "python/_autosummary/mlx.core.atleast_3d.rst", "python/_autosummary/mlx.core.broadcast_to.rst", "python/_autosummary/mlx.core.ceil.rst", "python/_autosummary/mlx.core.clip.rst", "python/_autosummary/mlx.core.compile.rst", "python/_autosummary/mlx.core.concatenate.rst", "python/_autosummary/mlx.core.conv1d.rst", "python/_autosummary/mlx.core.conv2d.rst", "python/_autosummary/mlx.core.conv_general.rst", "python/_autosummary/mlx.core.convolve.rst", "python/_autosummary/mlx.core.cos.rst", "python/_autosummary/mlx.core.cosh.rst", "python/_autosummary/mlx.core.cummax.rst", "python/_autosummary/mlx.core.cummin.rst", "python/_autosummary/mlx.core.cumprod.rst", "python/_autosummary/mlx.core.cumsum.rst", "python/_autosummary/mlx.core.default_device.rst", "python/_autosummary/mlx.core.default_stream.rst", "python/_autosummary/mlx.core.dequantize.rst", "python/_autosummary/mlx.core.diag.rst", "python/_autosummary/mlx.core.diagonal.rst", "python/_autosummary/mlx.core.disable_compile.rst", "python/_autosummary/mlx.core.divide.rst", "python/_autosummary/mlx.core.divmod.rst", "python/_autosummary/mlx.core.enable_compile.rst", "python/_autosummary/mlx.core.equal.rst", "python/_autosummary/mlx.core.erf.rst", "python/_autosummary/mlx.core.erfinv.rst", "python/_autosummary/mlx.core.eval.rst", "python/_autosummary/mlx.core.exp.rst", "python/_autosummary/mlx.core.expand_dims.rst", "python/_autosummary/mlx.core.expm1.rst", "python/_autosummary/mlx.core.eye.rst", "python/_autosummary/mlx.core.fast.layer_norm.rst", "python/_autosummary/mlx.core.fast.rms_norm.rst", "python/_autosummary/mlx.core.fast.rope.rst", "python/_autosummary/mlx.core.fast.scaled_dot_product_attention.rst", "python/_autosummary/mlx.core.fft.fft.rst", "python/_autosummary/mlx.core.fft.fft2.rst", "python/_autosummary/mlx.core.fft.fftn.rst", "python/_autosummary/mlx.core.fft.ifft.rst", "python/_autosummary/mlx.core.fft.ifft2.rst", "python/_autosummary/mlx.core.fft.ifftn.rst", "python/_autosummary/mlx.core.fft.irfft.rst", "python/_autosummary/mlx.core.fft.irfft2.rst", "python/_autosummary/mlx.core.fft.irfftn.rst", "python/_autosummary/mlx.core.fft.rfft.rst", "python/_autosummary/mlx.core.fft.rfft2.rst", "python/_autosummary/mlx.core.fft.rfftn.rst", "python/_autosummary/mlx.core.flatten.rst", "python/_autosummary/mlx.core.floor.rst", "python/_autosummary/mlx.core.floor_divide.rst", "python/_autosummary/mlx.core.full.rst", "python/_autosummary/mlx.core.grad.rst", "python/_autosummary/mlx.core.greater.rst", "python/_autosummary/mlx.core.greater_equal.rst", "python/_autosummary/mlx.core.identity.rst", "python/_autosummary/mlx.core.inner.rst", "python/_autosummary/mlx.core.isclose.rst", "python/_autosummary/mlx.core.isinf.rst", "python/_autosummary/mlx.core.isnan.rst", "python/_autosummary/mlx.core.isneginf.rst", "python/_autosummary/mlx.core.isposinf.rst", "python/_autosummary/mlx.core.issubdtype.rst", "python/_autosummary/mlx.core.jvp.rst", "python/_autosummary/mlx.core.less.rst", "python/_autosummary/mlx.core.less_equal.rst", "python/_autosummary/mlx.core.linalg.norm.rst", "python/_autosummary/mlx.core.linalg.qr.rst", "python/_autosummary/mlx.core.linspace.rst", "python/_autosummary/mlx.core.load.rst", "python/_autosummary/mlx.core.log.rst", "python/_autosummary/mlx.core.log10.rst", "python/_autosummary/mlx.core.log1p.rst", "python/_autosummary/mlx.core.log2.rst", "python/_autosummary/mlx.core.logaddexp.rst", "python/_autosummary/mlx.core.logical_and.rst", "python/_autosummary/mlx.core.logical_not.rst", "python/_autosummary/mlx.core.logical_or.rst", "python/_autosummary/mlx.core.logsumexp.rst", "python/_autosummary/mlx.core.matmul.rst", "python/_autosummary/mlx.core.max.rst", "python/_autosummary/mlx.core.maximum.rst", "python/_autosummary/mlx.core.mean.rst", "python/_autosummary/mlx.core.meshgrid.rst", "python/_autosummary/mlx.core.metal.get_active_memory.rst", "python/_autosummary/mlx.core.metal.get_cache_memory.rst", "python/_autosummary/mlx.core.metal.get_peak_memory.rst", "python/_autosummary/mlx.core.metal.is_available.rst", "python/_autosummary/mlx.core.metal.set_cache_limit.rst", "python/_autosummary/mlx.core.metal.set_memory_limit.rst", "python/_autosummary/mlx.core.metal.start_capture.rst", "python/_autosummary/mlx.core.metal.stop_capture.rst", "python/_autosummary/mlx.core.min.rst", "python/_autosummary/mlx.core.minimum.rst", "python/_autosummary/mlx.core.moveaxis.rst", "python/_autosummary/mlx.core.multiply.rst", "python/_autosummary/mlx.core.negative.rst", "python/_autosummary/mlx.core.new_stream.rst", "python/_autosummary/mlx.core.ones.rst", "python/_autosummary/mlx.core.ones_like.rst", "python/_autosummary/mlx.core.outer.rst", "python/_autosummary/mlx.core.pad.rst", "python/_autosummary/mlx.core.partition.rst", "python/_autosummary/mlx.core.prod.rst", "python/_autosummary/mlx.core.quantize.rst", "python/_autosummary/mlx.core.quantized_matmul.rst", "python/_autosummary/mlx.core.random.bernoulli.rst", "python/_autosummary/mlx.core.random.categorical.rst", "python/_autosummary/mlx.core.random.gumbel.rst", "python/_autosummary/mlx.core.random.key.rst", "python/_autosummary/mlx.core.random.multivariate_normal.rst", "python/_autosummary/mlx.core.random.normal.rst", "python/_autosummary/mlx.core.random.randint.rst", "python/_autosummary/mlx.core.random.seed.rst", "python/_autosummary/mlx.core.random.split.rst", "python/_autosummary/mlx.core.random.truncated_normal.rst", "python/_autosummary/mlx.core.random.uniform.rst", "python/_autosummary/mlx.core.reciprocal.rst", "python/_autosummary/mlx.core.repeat.rst", "python/_autosummary/mlx.core.reshape.rst", "python/_autosummary/mlx.core.round.rst", "python/_autosummary/mlx.core.rsqrt.rst", "python/_autosummary/mlx.core.save.rst", "python/_autosummary/mlx.core.save_gguf.rst", "python/_autosummary/mlx.core.save_safetensors.rst", "python/_autosummary/mlx.core.savez.rst", "python/_autosummary/mlx.core.savez_compressed.rst", "python/_autosummary/mlx.core.set_default_device.rst", "python/_autosummary/mlx.core.set_default_stream.rst", "python/_autosummary/mlx.core.sigmoid.rst", "python/_autosummary/mlx.core.sign.rst", "python/_autosummary/mlx.core.sin.rst", "python/_autosummary/mlx.core.sinh.rst", "python/_autosummary/mlx.core.softmax.rst", "python/_autosummary/mlx.core.sort.rst", "python/_autosummary/mlx.core.split.rst", "python/_autosummary/mlx.core.sqrt.rst", "python/_autosummary/mlx.core.square.rst", "python/_autosummary/mlx.core.squeeze.rst", "python/_autosummary/mlx.core.stack.rst", "python/_autosummary/mlx.core.std.rst", "python/_autosummary/mlx.core.stop_gradient.rst", "python/_autosummary/mlx.core.stream.rst", "python/_autosummary/mlx.core.subtract.rst", "python/_autosummary/mlx.core.sum.rst", "python/_autosummary/mlx.core.swapaxes.rst", "python/_autosummary/mlx.core.take.rst", "python/_autosummary/mlx.core.take_along_axis.rst", "python/_autosummary/mlx.core.tan.rst", "python/_autosummary/mlx.core.tanh.rst", "python/_autosummary/mlx.core.tensordot.rst", "python/_autosummary/mlx.core.tile.rst", "python/_autosummary/mlx.core.topk.rst", "python/_autosummary/mlx.core.transpose.rst", "python/_autosummary/mlx.core.tri.rst", "python/_autosummary/mlx.core.tril.rst", "python/_autosummary/mlx.core.triu.rst", "python/_autosummary/mlx.core.value_and_grad.rst", "python/_autosummary/mlx.core.var.rst", "python/_autosummary/mlx.core.vjp.rst", "python/_autosummary/mlx.core.vmap.rst", "python/_autosummary/mlx.core.where.rst", "python/_autosummary/mlx.core.zeros.rst", "python/_autosummary/mlx.core.zeros_like.rst", "python/_autosummary/mlx.nn.value_and_grad.rst", "python/_autosummary/mlx.utils.tree_flatten.rst", "python/_autosummary/mlx.utils.tree_map.rst", "python/_autosummary/mlx.utils.tree_unflatten.rst", "python/_autosummary/stream_class.rst", "python/array.rst", "python/data_types.rst", "python/devices_and_streams.rst", "python/fast.rst", "python/fft.rst", "python/linalg.rst", "python/metal.rst", "python/nn.rst", "python/nn/_autosummary/mlx.nn.ALiBi.rst", "python/nn/_autosummary/mlx.nn.AvgPool1d.rst", "python/nn/_autosummary/mlx.nn.AvgPool2d.rst", "python/nn/_autosummary/mlx.nn.BatchNorm.rst", "python/nn/_autosummary/mlx.nn.Conv1d.rst", "python/nn/_autosummary/mlx.nn.Conv2d.rst", "python/nn/_autosummary/mlx.nn.Dropout.rst", "python/nn/_autosummary/mlx.nn.Dropout2d.rst", "python/nn/_autosummary/mlx.nn.Dropout3d.rst", "python/nn/_autosummary/mlx.nn.Embedding.rst", "python/nn/_autosummary/mlx.nn.GELU.rst", "python/nn/_autosummary/mlx.nn.GRU.rst", "python/nn/_autosummary/mlx.nn.GroupNorm.rst", "python/nn/_autosummary/mlx.nn.InstanceNorm.rst", "python/nn/_autosummary/mlx.nn.LSTM.rst", "python/nn/_autosummary/mlx.nn.LayerNorm.rst", "python/nn/_autosummary/mlx.nn.Linear.rst", "python/nn/_autosummary/mlx.nn.MaxPool1d.rst", "python/nn/_autosummary/mlx.nn.MaxPool2d.rst", "python/nn/_autosummary/mlx.nn.Mish.rst", "python/nn/_autosummary/mlx.nn.Module.apply.rst", "python/nn/_autosummary/mlx.nn.Module.apply_to_modules.rst", "python/nn/_autosummary/mlx.nn.Module.children.rst", "python/nn/_autosummary/mlx.nn.Module.eval.rst", "python/nn/_autosummary/mlx.nn.Module.filter_and_map.rst", "python/nn/_autosummary/mlx.nn.Module.freeze.rst", "python/nn/_autosummary/mlx.nn.Module.leaf_modules.rst", "python/nn/_autosummary/mlx.nn.Module.load_weights.rst", "python/nn/_autosummary/mlx.nn.Module.modules.rst", "python/nn/_autosummary/mlx.nn.Module.named_modules.rst", "python/nn/_autosummary/mlx.nn.Module.parameters.rst", "python/nn/_autosummary/mlx.nn.Module.save_weights.rst", "python/nn/_autosummary/mlx.nn.Module.set_dtype.rst", "python/nn/_autosummary/mlx.nn.Module.state.rst", "python/nn/_autosummary/mlx.nn.Module.train.rst", "python/nn/_autosummary/mlx.nn.Module.trainable_parameters.rst", "python/nn/_autosummary/mlx.nn.Module.training.rst", "python/nn/_autosummary/mlx.nn.Module.unfreeze.rst", "python/nn/_autosummary/mlx.nn.Module.update.rst", "python/nn/_autosummary/mlx.nn.Module.update_modules.rst", "python/nn/_autosummary/mlx.nn.MultiHeadAttention.rst", "python/nn/_autosummary/mlx.nn.PReLU.rst", "python/nn/_autosummary/mlx.nn.QuantizedLinear.rst", "python/nn/_autosummary/mlx.nn.RMSNorm.rst", "python/nn/_autosummary/mlx.nn.RNN.rst", "python/nn/_autosummary/mlx.nn.ReLU.rst", "python/nn/_autosummary/mlx.nn.RoPE.rst", "python/nn/_autosummary/mlx.nn.SELU.rst", "python/nn/_autosummary/mlx.nn.Sequential.rst", "python/nn/_autosummary/mlx.nn.SiLU.rst", "python/nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.rst", "python/nn/_autosummary/mlx.nn.Softshrink.rst", "python/nn/_autosummary/mlx.nn.Step.rst", "python/nn/_autosummary/mlx.nn.Transformer.rst", "python/nn/_autosummary/mlx.nn.Upsample.rst", "python/nn/_autosummary/mlx.nn.init.constant.rst", "python/nn/_autosummary/mlx.nn.init.glorot_normal.rst", "python/nn/_autosummary/mlx.nn.init.glorot_uniform.rst", "python/nn/_autosummary/mlx.nn.init.he_normal.rst", "python/nn/_autosummary/mlx.nn.init.he_uniform.rst", "python/nn/_autosummary/mlx.nn.init.identity.rst", "python/nn/_autosummary/mlx.nn.init.normal.rst", "python/nn/_autosummary/mlx.nn.init.uniform.rst", "python/nn/_autosummary_functions/mlx.nn.elu.rst", "python/nn/_autosummary_functions/mlx.nn.gelu.rst", "python/nn/_autosummary_functions/mlx.nn.gelu_approx.rst", "python/nn/_autosummary_functions/mlx.nn.gelu_fast_approx.rst", "python/nn/_autosummary_functions/mlx.nn.glu.rst", "python/nn/_autosummary_functions/mlx.nn.hardswish.rst", "python/nn/_autosummary_functions/mlx.nn.leaky_relu.rst", "python/nn/_autosummary_functions/mlx.nn.log_sigmoid.rst", "python/nn/_autosummary_functions/mlx.nn.log_softmax.rst", "python/nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.rst", "python/nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.cross_entropy.rst", "python/nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.hinge_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.huber_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.l1_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.mse_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.nll_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.rst", "python/nn/_autosummary_functions/mlx.nn.losses.triplet_loss.rst", "python/nn/_autosummary_functions/mlx.nn.mish.rst", "python/nn/_autosummary_functions/mlx.nn.prelu.rst", "python/nn/_autosummary_functions/mlx.nn.relu.rst", "python/nn/_autosummary_functions/mlx.nn.relu6.rst", "python/nn/_autosummary_functions/mlx.nn.selu.rst", "python/nn/_autosummary_functions/mlx.nn.sigmoid.rst", "python/nn/_autosummary_functions/mlx.nn.silu.rst", "python/nn/_autosummary_functions/mlx.nn.softmax.rst", "python/nn/_autosummary_functions/mlx.nn.softplus.rst", "python/nn/_autosummary_functions/mlx.nn.softshrink.rst", "python/nn/_autosummary_functions/mlx.nn.step.rst", "python/nn/_autosummary_functions/mlx.nn.tanh.rst", "python/nn/functions.rst", "python/nn/init.rst", "python/nn/layers.rst", "python/nn/losses.rst", "python/nn/module.rst", "python/ops.rst", "python/optimizers.rst", "python/optimizers/_autosummary/mlx.optimizers.AdaDelta.rst", "python/optimizers/_autosummary/mlx.optimizers.Adafactor.rst", "python/optimizers/_autosummary/mlx.optimizers.Adagrad.rst", "python/optimizers/_autosummary/mlx.optimizers.Adam.rst", "python/optimizers/_autosummary/mlx.optimizers.AdamW.rst", "python/optimizers/_autosummary/mlx.optimizers.Adamax.rst", "python/optimizers/_autosummary/mlx.optimizers.Lion.rst", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.apply_gradients.rst", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.init.rst", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.state.rst", "python/optimizers/_autosummary/mlx.optimizers.Optimizer.update.rst", "python/optimizers/_autosummary/mlx.optimizers.RMSprop.rst", "python/optimizers/_autosummary/mlx.optimizers.SGD.rst", "python/optimizers/_autosummary/mlx.optimizers.cosine_decay.rst", "python/optimizers/_autosummary/mlx.optimizers.exponential_decay.rst", "python/optimizers/_autosummary/mlx.optimizers.join_schedules.rst", "python/optimizers/_autosummary/mlx.optimizers.linear_schedule.rst", "python/optimizers/_autosummary/mlx.optimizers.step_decay.rst", "python/optimizers/common_optimizers.rst", "python/optimizers/optimizer.rst", "python/optimizers/schedulers.rst", "python/random.rst", "python/transforms.rst", "python/tree_utils.rst", "usage/compile.rst", "usage/function_transforms.rst", "usage/indexing.rst", "usage/lazy_evaluation.rst", "usage/numpy.rst", "usage/quick_start.rst", "usage/saving_and_loading.rst", "usage/unified_memory.rst", "usage/using_streams.rst"], "titles": ["Operations", "Developer Documentation", "Metal Debugger", "Linear Regression", "LLM inference", "Multi-Layer Perceptron", "MLX", "Build and Install", "mlx.core.Device", "mlx.core.Dtype", "mlx.core.DtypeCategory", "mlx.core.abs", "mlx.core.add", "mlx.core.all", "mlx.core.allclose", "mlx.core.any", "mlx.core.arange", "mlx.core.arccos", "mlx.core.arccosh", "mlx.core.arcsin", "mlx.core.arcsinh", "mlx.core.arctan", "mlx.core.arctanh", "mlx.core.argmax", "mlx.core.argmin", "mlx.core.argpartition", "mlx.core.argsort", "mlx.core.array", "mlx.core.array.T", "mlx.core.array.abs", "mlx.core.array.all", "mlx.core.array.any", "mlx.core.array.argmax", "mlx.core.array.argmin", "mlx.core.array.astype", "mlx.core.array.at", "mlx.core.array.cos", "mlx.core.array.cummax", "mlx.core.array.cummin", "mlx.core.array.cumprod", "mlx.core.array.cumsum", "mlx.core.array.diag", "mlx.core.array.diagonal", "mlx.core.array.dtype", "mlx.core.array.exp", "mlx.core.array.flatten", "mlx.core.array.item", "mlx.core.array.itemsize", "mlx.core.array.log", "mlx.core.array.log10", "mlx.core.array.log1p", "mlx.core.array.log2", "mlx.core.array.logsumexp", "mlx.core.array.max", "mlx.core.array.mean", "mlx.core.array.min", "mlx.core.array.moveaxis", "mlx.core.array.nbytes", "mlx.core.array.ndim", "mlx.core.array.prod", "mlx.core.array.reciprocal", "mlx.core.array.reshape", "mlx.core.array.round", "mlx.core.array.rsqrt", "mlx.core.array.shape", "mlx.core.array.sin", "mlx.core.array.size", "mlx.core.array.split", "mlx.core.array.sqrt", "mlx.core.array.square", "mlx.core.array.squeeze", "mlx.core.array.sum", "mlx.core.array.swapaxes", "mlx.core.array.tolist", "mlx.core.array.transpose", "mlx.core.array.var", "mlx.core.array_equal", "mlx.core.atleast_1d", "mlx.core.atleast_2d", "mlx.core.atleast_3d", "mlx.core.broadcast_to", "mlx.core.ceil", "mlx.core.clip", "mlx.core.compile", "mlx.core.concatenate", "mlx.core.conv1d", "mlx.core.conv2d", "mlx.core.conv_general", "mlx.core.convolve", "mlx.core.cos", "mlx.core.cosh", "mlx.core.cummax", "mlx.core.cummin", "mlx.core.cumprod", "mlx.core.cumsum", "mlx.core.default_device", "mlx.core.default_stream", "mlx.core.dequantize", "mlx.core.diag", "mlx.core.diagonal", "mlx.core.disable_compile", "mlx.core.divide", "mlx.core.divmod", "mlx.core.enable_compile", "mlx.core.equal", "mlx.core.erf", "mlx.core.erfinv", "mlx.core.eval", "mlx.core.exp", "mlx.core.expand_dims", "mlx.core.expm1", "mlx.core.eye", "mlx.core.fast.layer_norm", "mlx.core.fast.rms_norm", "mlx.core.fast.rope", "mlx.core.fast.scaled_dot_product_attention", "mlx.core.fft.fft", "mlx.core.fft.fft2", "mlx.core.fft.fftn", "mlx.core.fft.ifft", "mlx.core.fft.ifft2", "mlx.core.fft.ifftn", "mlx.core.fft.irfft", "mlx.core.fft.irfft2", "mlx.core.fft.irfftn", "mlx.core.fft.rfft", "mlx.core.fft.rfft2", "mlx.core.fft.rfftn", "mlx.core.flatten", "mlx.core.floor", "mlx.core.floor_divide", "mlx.core.full", "mlx.core.grad", "mlx.core.greater", "mlx.core.greater_equal", "mlx.core.identity", "mlx.core.inner", "mlx.core.isclose", "mlx.core.isinf", "mlx.core.isnan", "mlx.core.isneginf", "mlx.core.isposinf", "mlx.core.issubdtype", "mlx.core.jvp", "mlx.core.less", "mlx.core.less_equal", "mlx.core.linalg.norm", "mlx.core.linalg.qr", "mlx.core.linspace", "mlx.core.load", "mlx.core.log", "mlx.core.log10", "mlx.core.log1p", "mlx.core.log2", "mlx.core.logaddexp", "mlx.core.logical_and", "mlx.core.logical_not", "mlx.core.logical_or", "mlx.core.logsumexp", "mlx.core.matmul", "mlx.core.max", "mlx.core.maximum", "mlx.core.mean", "mlx.core.meshgrid", "mlx.core.metal.get_active_memory", "mlx.core.metal.get_cache_memory", "mlx.core.metal.get_peak_memory", "mlx.core.metal.is_available", "mlx.core.metal.set_cache_limit", "mlx.core.metal.set_memory_limit", "mlx.core.metal.start_capture", "mlx.core.metal.stop_capture", "mlx.core.min", "mlx.core.minimum", "mlx.core.moveaxis", "mlx.core.multiply", "mlx.core.negative", "mlx.core.new_stream", "mlx.core.ones", "mlx.core.ones_like", "mlx.core.outer", "mlx.core.pad", "mlx.core.partition", "mlx.core.prod", "mlx.core.quantize", "mlx.core.quantized_matmul", "mlx.core.random.bernoulli", "mlx.core.random.categorical", "mlx.core.random.gumbel", "mlx.core.random.key", "mlx.core.random.multivariate_normal", "mlx.core.random.normal", "mlx.core.random.randint", "mlx.core.random.seed", "mlx.core.random.split", "mlx.core.random.truncated_normal", "mlx.core.random.uniform", "mlx.core.reciprocal", "mlx.core.repeat", "mlx.core.reshape", "mlx.core.round", "mlx.core.rsqrt", "mlx.core.save", "mlx.core.save_gguf", "mlx.core.save_safetensors", "mlx.core.savez", "mlx.core.savez_compressed", "mlx.core.set_default_device", "mlx.core.set_default_stream", "mlx.core.sigmoid", "mlx.core.sign", "mlx.core.sin", "mlx.core.sinh", "mlx.core.softmax", "mlx.core.sort", "mlx.core.split", "mlx.core.sqrt", "mlx.core.square", "mlx.core.squeeze", "mlx.core.stack", "mlx.core.std", "mlx.core.stop_gradient", "mlx.core.stream", "mlx.core.subtract", "mlx.core.sum", "mlx.core.swapaxes", "mlx.core.take", "mlx.core.take_along_axis", "mlx.core.tan", "mlx.core.tanh", "mlx.core.tensordot", "mlx.core.tile", "mlx.core.topk", "mlx.core.transpose", "mlx.core.tri", "mlx.core.tril", "mlx.core.triu", "mlx.core.value_and_grad", "mlx.core.var", "mlx.core.vjp", "mlx.core.vmap", "mlx.core.where", "mlx.core.zeros", "mlx.core.zeros_like", "mlx.nn.value_and_grad", "mlx.utils.tree_flatten", "mlx.utils.tree_map", "mlx.utils.tree_unflatten", "mlx.core.Stream", "Array", "Data Types", "Devices and Streams", "Fast", "FFT", "Linear Algebra", "Metal", "Neural Networks", "mlx.nn.ALiBi", "mlx.nn.AvgPool1d", "mlx.nn.AvgPool2d", "mlx.nn.BatchNorm", "mlx.nn.Conv1d", "mlx.nn.Conv2d", "mlx.nn.Dropout", "mlx.nn.Dropout2d", "mlx.nn.Dropout3d", "mlx.nn.Embedding", "mlx.nn.GELU", "mlx.nn.GRU", "mlx.nn.GroupNorm", "mlx.nn.InstanceNorm", "mlx.nn.LSTM", "mlx.nn.LayerNorm", "mlx.nn.Linear", "mlx.nn.MaxPool1d", "mlx.nn.MaxPool2d", "mlx.nn.Mish", "mlx.nn.Module.apply", "mlx.nn.Module.apply_to_modules", "mlx.nn.Module.children", "mlx.nn.Module.eval", "mlx.nn.Module.filter_and_map", "mlx.nn.Module.freeze", "mlx.nn.Module.leaf_modules", "mlx.nn.Module.load_weights", "mlx.nn.Module.modules", "mlx.nn.Module.named_modules", "mlx.nn.Module.parameters", "mlx.nn.Module.save_weights", "mlx.nn.Module.set_dtype", "mlx.nn.Module.state", "mlx.nn.Module.train", "mlx.nn.Module.trainable_parameters", "mlx.nn.Module.training", "mlx.nn.Module.unfreeze", "mlx.nn.Module.update", "mlx.nn.Module.update_modules", "mlx.nn.MultiHeadAttention", "mlx.nn.PReLU", "mlx.nn.QuantizedLinear", "mlx.nn.RMSNorm", "mlx.nn.RNN", "mlx.nn.ReLU", "mlx.nn.RoPE", "mlx.nn.SELU", "mlx.nn.Sequential", "mlx.nn.SiLU", "mlx.nn.SinusoidalPositionalEncoding", "mlx.nn.Softshrink", "mlx.nn.Step", "mlx.nn.Transformer", "mlx.nn.Upsample", "mlx.nn.init.constant", "mlx.nn.init.glorot_normal", "mlx.nn.init.glorot_uniform", "mlx.nn.init.he_normal", "mlx.nn.init.he_uniform", "mlx.nn.init.identity", "mlx.nn.init.normal", "mlx.nn.init.uniform", "mlx.nn.elu", "mlx.nn.gelu", "mlx.nn.gelu_approx", "mlx.nn.gelu_fast_approx", "mlx.nn.glu", "mlx.nn.hardswish", "mlx.nn.leaky_relu", "mlx.nn.log_sigmoid", "mlx.nn.log_softmax", "mlx.nn.losses.binary_cross_entropy", "mlx.nn.losses.cosine_similarity_loss", "mlx.nn.losses.cross_entropy", "mlx.nn.losses.gaussian_nll_loss", "mlx.nn.losses.hinge_loss", "mlx.nn.losses.huber_loss", "mlx.nn.losses.kl_div_loss", "mlx.nn.losses.l1_loss", "mlx.nn.losses.log_cosh_loss", "mlx.nn.losses.margin_ranking_loss", "mlx.nn.losses.mse_loss", "mlx.nn.losses.nll_loss", "mlx.nn.losses.smooth_l1_loss", "mlx.nn.losses.triplet_loss", "mlx.nn.mish", "mlx.nn.prelu", "mlx.nn.relu", "mlx.nn.relu6", "mlx.nn.selu", "mlx.nn.sigmoid", "mlx.nn.silu", "mlx.nn.softmax", "mlx.nn.softplus", "mlx.nn.softshrink", "mlx.nn.step", "mlx.nn.tanh", "Functions", "Initializers", "Layers", "Loss Functions", "Module", "Operations", "Optimizers", "mlx.optimizers.AdaDelta", "mlx.optimizers.Adafactor", "mlx.optimizers.Adagrad", "mlx.optimizers.Adam", "mlx.optimizers.AdamW", "mlx.optimizers.Adamax", "mlx.optimizers.Lion", "mlx.optimizers.Optimizer.apply_gradients", "mlx.optimizers.Optimizer.init", "mlx.optimizers.Optimizer.state", "mlx.optimizers.Optimizer.update", "mlx.optimizers.RMSprop", "mlx.optimizers.SGD", "mlx.optimizers.cosine_decay", "mlx.optimizers.exponential_decay", "mlx.optimizers.join_schedules", "mlx.optimizers.linear_schedule", "mlx.optimizers.step_decay", "Common Optimizers", "Optimizer", "Schedulers", "Random", "Transforms", "Tree Utils", "Compilation", "Function Transforms", "Indexing Arrays", "Lazy Evaluation", "Conversion to NumPy and Other Frameworks", "Quick Start Guide", "Saving and Loading Arrays", "Unified Memory", "Using Streams"], "terms": {"mlx": [1, 2, 3, 4, 5, 7, 256, 356, 359, 361, 383, 385, 386, 387, 388, 389, 390, 391, 392, 393], "provid": [1, 4, 97, 132, 230, 237, 246, 256, 277, 282, 284, 294, 295, 296, 299, 310, 311, 355, 359, 392, 394], "open": [2, 7, 16, 192, 196], "flexibl": 6, "which": [1, 4, 5, 6, 7, 16, 34, 83, 87, 99, 107, 114, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 132, 138, 139, 140, 141, 143, 146, 147, 149, 163, 170, 184, 187, 188, 198, 199, 202, 203, 204, 205, 206, 218, 219, 226, 237, 239, 240, 259, 264, 265, 267, 275, 277, 281, 303, 331, 334, 338, 341, 356, 369, 370, 383, 386, 387, 388, 389, 393, 394], "user": [1, 4, 256], "mai": [1, 146, 264, 387, 388], "add": [1, 2, 4, 35, 109, 154, 181, 184, 261, 262, 387, 393], "special": 1, "without": [4, 6, 221, 297, 355, 385, 386, 389, 390, 393], "much": [1, 4, 258, 259, 274, 275, 386, 389], "hassl": [], "while": [1, 2, 4, 7, 199, 303, 389, 390], "librari": [1, 7, 256], "suppli": [], "effici": [4, 6, 264, 303, 389, 391], "can": [1, 2, 4, 6, 7, 12, 16, 61, 74, 83, 99, 100, 101, 102, 104, 107, 133, 134, 144, 145, 146, 154, 161, 173, 175, 186, 187, 192, 195, 196, 203, 223, 237, 256, 259, 266, 275, 281, 294, 305, 311, 331, 356, 359, 361, 369, 370, 383, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394], "compos": [6, 256, 386, 387, 391], "ani": [1, 4, 6, 16, 83, 245, 246, 247, 256, 267, 277, 278, 281, 290, 299, 310, 311, 356, 378, 385, 386, 387, 389, 391, 392, 393], "number": [1, 10, 16, 57, 66, 83, 86, 87, 97, 111, 132, 135, 143, 148, 181, 184, 185, 187, 191, 194, 196, 198, 200, 230, 231, 234, 237, 239, 240, 256, 260, 261, 262, 264, 265, 269, 270, 297, 298, 310, 311, 313, 314, 315, 316, 375, 377, 378, 383, 386, 387, 394], "applic": [2, 7], "aris": 390, "case": [1, 4, 118, 121, 122, 124, 125, 126, 127, 128, 147, 159, 199, 218, 259, 264, 275, 309, 341, 347, 352, 353, 369, 370, 386, 387, 391, 392, 393, 394], "where": [5, 111, 137, 184, 237, 240, 258, 259, 260, 261, 262, 263, 264, 265, 267, 268, 269, 270, 271, 272, 273, 274, 275, 281, 298, 300, 301, 309, 315, 316, 320, 321, 322, 323, 332, 338, 344, 347, 349, 353, 370, 387, 388], "new": [1, 5, 80, 99, 174, 177, 199, 219, 233, 246, 289, 297, 359, 361, 372, 377, 386, 388, 389, 390], "function": [1, 2, 3, 4, 5, 6, 14, 83, 102, 105, 106, 132, 137, 143, 146, 147, 159, 209, 237, 239, 240, 244, 246, 256, 267, 276, 278, 282, 289, 294, 298, 301, 302, 304, 305, 306, 308, 309, 310, 321, 322, 323, 324, 325, 327, 328, 343, 348, 350, 351, 352, 353, 354, 356, 361, 370, 383, 385, 388, 389, 390, 392], "highli": 7, "optim": [2, 3, 5, 6, 295, 386, 387, 389], "ar": [1, 3, 4, 5, 6, 7, 14, 16, 76, 80, 82, 83, 87, 88, 99, 107, 111, 117, 118, 120, 121, 123, 124, 126, 127, 128, 132, 137, 138, 139, 140, 141, 142, 143, 146, 147, 149, 159, 169, 180, 181, 182, 184, 185, 186, 187, 188, 192, 195, 196, 205, 206, 218, 219, 226, 237, 239, 240, 245, 246, 250, 260, 261, 262, 263, 264, 265, 269, 270, 272, 273, 284, 297, 299, 311, 329, 331, 332, 355, 359, 368, 370, 385, 386, 387, 388, 389, 390, 391, 392, 393], "need": [1, 4, 5, 6, 76, 184, 256, 295, 296, 307, 310, 383, 387, 389, 390, 391, 393], "For": [1, 4, 7, 35, 115, 142, 146, 184, 247, 256, 260, 264, 277, 282, 291, 294, 299, 303, 307, 311, 313, 314, 315, 316, 356, 383, 386, 387, 388, 389, 390, 391, 392, 393], "you": [1, 2, 4, 5, 6, 7, 256, 307, 310, 356, 383, 386, 387, 388, 390, 392, 393], "design": [3, 6, 383, 393], "your": [1, 4, 7, 359, 387, 389], "own": [7, 390], "link": [1, 7], "top": [1, 232, 273, 311], "core": [1, 2, 3, 4, 5, 256, 258, 259, 260, 270, 274, 275, 284, 287, 289, 292, 311, 312, 313, 314, 315, 316, 317, 318, 319, 329, 331, 338, 356, 359, 361, 386, 390, 391], "we": [1, 3, 4, 5, 97, 184, 185, 256, 266, 305, 366, 368, 383, 385, 386, 387, 389, 393], "inner": 386, "work": [1, 2, 4, 7, 169, 386, 387, 388, 389], "go": [1, 4, 387], "over": [1, 4, 5, 13, 15, 23, 24, 25, 26, 85, 86, 87, 91, 92, 93, 94, 118, 121, 124, 127, 136, 146, 148, 158, 160, 162, 172, 182, 183, 201, 213, 214, 220, 224, 230, 232, 238, 260, 261, 262, 269, 272, 300, 331, 375, 378, 387], "simpl": [1, 4, 5, 256, 266, 355, 386, 387, 389], "learn": [3, 5, 6, 260, 269, 270, 272, 298, 300, 362, 363, 364, 365, 366, 367, 368, 373, 374], "step": [2, 4, 5, 16, 256, 268, 271, 301, 363, 370, 375, 377, 378, 379, 386], "involv": [361, 386], "ad": [1, 3, 7, 112, 270, 359, 362, 363, 364, 365, 366, 367, 373, 389, 392], "let": [1, 3, 4, 386, 387, 389, 390], "": [1, 3, 4, 5, 43, 47, 58, 83, 96, 97, 117, 118, 120, 121, 123, 124, 126, 127, 132, 146, 149, 162, 180, 184, 187, 200, 203, 204, 220, 222, 237, 238, 240, 244, 256, 259, 268, 271, 275, 281, 282, 284, 288, 289, 290, 294, 301, 361, 370, 371, 383, 386, 387, 389, 390, 391, 392, 393], "sai": [1, 4, 356, 389], "would": [1, 4, 311, 388, 389, 390, 393], "like": [1, 4, 6, 142, 179, 243, 265, 337, 370, 372, 386, 387, 389, 390, 391, 393], "an": [1, 2, 4, 5, 7, 9, 13, 15, 27, 77, 78, 79, 80, 85, 86, 87, 107, 111, 112, 115, 128, 131, 135, 146, 149, 169, 174, 178, 179, 181, 183, 184, 185, 198, 199, 200, 215, 218, 225, 226, 227, 230, 231, 234, 240, 242, 243, 245, 246, 256, 258, 259, 263, 269, 271, 272, 273, 274, 275, 277, 297, 298, 299, 301, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 322, 344, 356, 362, 372, 376, 381, 383, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394], "take": [1, 4, 5, 83, 132, 143, 161, 173, 179, 185, 227, 237, 239, 240, 243, 297, 383, 387, 388, 392, 393, 394], "two": [1, 12, 14, 76, 78, 99, 101, 104, 117, 120, 126, 133, 134, 137, 144, 145, 147, 154, 159, 161, 173, 175, 180, 225, 259, 271, 275, 299, 324, 330, 386, 387, 388, 393], "arrai": [1, 4, 5, 6, 9, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 97, 98, 99, 101, 102, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 172, 173, 174, 175, 176, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 256, 260, 271, 277, 284, 287, 292, 298, 311, 312, 313, 314, 315, 316, 317, 318, 319, 321, 324, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 353, 356, 359, 362, 363, 364, 365, 366, 367, 368, 373, 374, 375, 376, 377, 378, 379, 386, 387, 389, 390, 391, 393], "x": [1, 3, 4, 5, 35, 105, 110, 112, 113, 135, 146, 185, 188, 200, 205, 209, 235, 236, 241, 246, 256, 258, 259, 260, 267, 269, 270, 272, 273, 274, 275, 276, 277, 298, 300, 302, 307, 309, 311, 320, 321, 322, 323, 324, 325, 326, 327, 328, 341, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 359, 361, 368, 386, 387, 388, 389, 390, 391, 393], "y": [1, 3, 4, 5, 35, 241, 256, 260, 264, 269, 270, 272, 273, 300, 333, 338, 341, 361, 364, 386, 387, 389, 390], "scale": [1, 4, 97, 112, 113, 114, 115, 184, 185, 191, 264, 265, 272, 297, 303, 304, 307, 311, 347, 363], "them": [1, 4, 256, 282, 294, 393], "both": [1, 12, 101, 102, 104, 133, 134, 142, 144, 145, 146, 154, 161, 173, 175, 187, 223, 258, 259, 270, 271, 274, 275, 361, 386, 387, 391, 393], "some": [1, 3, 4, 5, 282, 294, 370, 386, 387, 389], "coeffici": [1, 362, 363, 365, 366, 367, 368], "alpha": [1, 184, 320, 342, 344, 347, 366, 373], "beta": [1, 97, 184, 260, 269, 270, 272, 341, 365, 366, 367, 368], "respect": [1, 3, 5, 112, 113, 132, 184, 237, 246, 256, 260, 267, 269, 270, 272, 359, 387, 391], "togeth": [1, 5, 184, 246], "get": [1, 3, 5, 7, 86, 87, 95, 96, 164, 165, 166, 189, 256, 386, 387, 389, 393], "z": [1, 268, 386, 389], "well": [4, 256, 282, 294, 297, 389], "veri": [4, 297, 389, 393], "easili": [], "do": [1, 4, 7, 256, 283, 294, 356, 359, 366, 386, 387, 389], "just": [1, 5, 272, 386, 388], "write": [1, 4, 256, 390], "out": [1, 7, 258, 259, 264, 265, 274, 275, 291, 386, 387, 388], "follow": [1, 4, 5, 6, 7, 16, 88, 97, 146, 184, 256, 322, 323, 335, 362, 363, 364, 365, 366, 367, 368, 374, 383, 386, 387, 393], "import": [1, 2, 3, 4, 5, 7, 146, 205, 237, 245, 246, 247, 256, 258, 259, 260, 270, 274, 275, 284, 311, 329, 331, 338, 356, 359, 386, 387, 388, 389, 390, 391], "mx": [1, 2, 3, 4, 5, 35, 128, 142, 146, 147, 149, 205, 237, 256, 258, 259, 260, 270, 274, 275, 277, 284, 288, 302, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 326, 329, 330, 331, 335, 338, 345, 354, 356, 359, 361, 383, 386, 387, 388, 389, 390, 391, 392, 393, 394], "def": [1, 3, 4, 5, 237, 256, 359, 386, 387, 388, 389, 390, 393], "simple_axpbi": 1, "float": [1, 10, 14, 16, 73, 112, 113, 114, 115, 130, 131, 137, 142, 146, 185, 186, 191, 250, 260, 263, 264, 265, 269, 270, 272, 277, 289, 300, 303, 307, 309, 310, 311, 312, 313, 314, 315, 316, 318, 319, 330, 331, 332, 334, 338, 341, 342, 352, 353, 362, 363, 364, 365, 366, 367, 368, 373, 374, 375, 376, 378, 379], "return": [1, 3, 4, 5, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 34, 46, 64, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 97, 98, 99, 101, 102, 104, 105, 106, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 165, 168, 169, 170, 172, 173, 174, 175, 176, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 194, 195, 196, 197, 198, 199, 200, 201, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 256, 268, 271, 277, 278, 279, 281, 282, 283, 284, 285, 286, 287, 291, 292, 294, 295, 296, 299, 301, 312, 313, 314, 315, 316, 317, 318, 319, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 356, 359, 369, 385, 386, 387, 388, 389, 390, 392, 393], "thi": [1, 4, 5, 7, 13, 14, 15, 16, 23, 24, 25, 26, 103, 137, 143, 146, 147, 154, 158, 159, 160, 162, 164, 172, 182, 183, 187, 208, 213, 214, 215, 220, 224, 226, 232, 238, 256, 263, 264, 265, 268, 271, 278, 279, 281, 282, 285, 286, 287, 292, 294, 295, 296, 297, 299, 301, 309, 313, 314, 315, 316, 322, 323, 324, 337, 353, 359, 370, 385, 386, 387, 389, 390, 392], "perform": [1, 2, 4, 6, 87, 91, 92, 93, 94, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 159, 185, 200, 213, 226, 256, 269, 310, 315, 316, 386, 388, 389, 393], "leav": [1, 107, 246], "differenti": [1, 6], "howev": [1, 256, 267, 269, 370, 383, 386, 389, 390], "vector": [1, 3, 6, 136, 143, 146, 226, 239, 240, 266, 331, 391], "math": [4, 342, 386], "often": 265, "realiz": [], "axpbi": 1, "routin": 1, "defin": [1, 3, 4, 5, 7, 146, 185, 245, 390], "same": [1, 4, 7, 14, 35, 76, 80, 83, 86, 87, 88, 112, 113, 122, 125, 126, 127, 132, 137, 143, 181, 187, 200, 239, 241, 256, 259, 260, 263, 269, 270, 275, 299, 312, 313, 314, 315, 316, 317, 318, 319, 331, 342, 359, 369, 383, 386, 388, 393], "realli": 272, "part": [1, 387, 388], "doe": [1, 2, 4, 7, 164, 256, 386, 388, 389, 390], "fast": [6, 267, 323, 393], "so": [1, 4, 7, 132, 237, 263, 311, 361, 386, 389, 393], "decid": [246, 281], "want": [4, 387, 393], "reli": 1, "acceler": [1, 260], "framework": [1, 6], "continu": 387, "impos": [], "our": [1, 4, 5, 305, 362, 363, 364, 365, 367, 368], "assumpt": [], "also": [1, 4, 5, 6, 7, 10, 12, 100, 101, 102, 104, 118, 121, 124, 127, 133, 134, 144, 145, 154, 161, 173, 175, 184, 223, 244, 256, 281, 295, 297, 299, 306, 321, 347, 349, 355, 361, 386, 387, 388, 389, 390, 391, 394], "assum": [1, 4, 147, 246, 256, 258, 259, 269, 274, 275], "how": [1, 4, 5, 256, 258, 259, 261, 262, 266, 274, 275, 311, 369, 386, 388, 393], "gradient": [3, 5, 132, 221, 237, 244, 256, 282, 295, 299, 310, 337, 359, 361, 362, 363, 365, 366, 367, 368, 369, 372, 374, 386, 387, 388, 389, 390, 391], "ins": [], "what": [1, 4, 246], "coincid": [], "right": [1, 7, 184, 258, 259, 267, 274, 275, 311, 322, 323, 332, 334, 342], "place": [1, 4, 35, 200, 389, 390], "cours": 387, "The": [1, 2, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 34, 43, 47, 57, 58, 64, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 97, 98, 99, 101, 102, 104, 105, 106, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 133, 134, 135, 136, 137, 138, 139, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 165, 166, 168, 169, 170, 172, 173, 174, 175, 176, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 194, 195, 196, 197, 198, 199, 203, 204, 209, 210, 211, 212, 213, 214, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 250, 258, 259, 260, 261, 262, 263, 264, 265, 266, 268, 269, 270, 271, 272, 273, 274, 275, 277, 278, 282, 284, 288, 289, 290, 291, 294, 295, 296, 297, 299, 300, 301, 303, 305, 307, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 324, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 353, 356, 359, 361, 362, 363, 364, 365, 366, 367, 368, 371, 373, 374, 375, 378, 381, 386, 387, 388, 389, 390, 391, 392, 393, 394], "structur": [1, 369, 387], "from": [1, 4, 5, 6, 97, 99, 123, 124, 126, 127, 131, 146, 149, 159, 163, 166, 168, 179, 184, 186, 187, 188, 189, 192, 195, 205, 218, 221, 223, 226, 227, 232, 241, 243, 245, 246, 247, 256, 273, 282, 284, 297, 313, 314, 315, 316, 318, 319, 332, 341, 356, 385, 386, 387, 389, 390, 391, 392, 393], "frontend": [], "api": [1, 387], "redirect": 1, "when": [1, 4, 6, 7, 83, 87, 146, 149, 261, 262, 311, 315, 316, 335, 341, 359, 377, 383, 386, 393], "appropri": [1, 386], "fallback": 1, "metal": [1, 6], "vjp": [1, 391], "jvp": [1, 391], "In": [1, 4, 5, 35, 159, 184, 246, 256, 264, 269, 359, 362, 364, 365, 367, 368, 369, 385, 386, 387, 389, 392, 393], "one": [1, 4, 7, 35, 73, 77, 82, 86, 87, 109, 111, 112, 113, 146, 152, 159, 185, 187, 218, 223, 250, 294, 311, 331, 393], "sentenc": [], "comput": [1, 3, 4, 5, 6, 7, 91, 92, 93, 94, 97, 110, 114, 132, 143, 146, 154, 162, 180, 184, 213, 220, 221, 230, 237, 238, 239, 244, 256, 260, 268, 269, 270, 271, 272, 282, 295, 299, 300, 303, 310, 313, 314, 315, 316, 322, 323, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 361, 362, 363, 365, 366, 367, 368, 372, 386, 387, 391, 393], "graph": [1, 4, 5, 6, 387], "rule": 1, "evalu": [1, 4, 5, 6, 107, 143, 239, 256, 280, 291, 359, 361, 386, 391], "said": 4, "start": [1, 3, 4, 6, 7, 16, 114, 148, 170, 215, 386, 388, 393], "discuss": 1, "more": [1, 2, 5, 9, 73, 99, 159, 168, 169, 203, 204, 250, 256, 260, 264, 303, 307, 310, 311, 313, 314, 315, 316, 383, 386, 387, 388, 391, 393], "detail": [1, 9, 168, 256, 264, 303, 307, 311, 313, 314, 315, 316, 362, 364, 365, 367, 368, 388, 391], "thei": [1, 3, 4, 14, 88, 137, 305, 333, 359, 368, 385, 386, 389, 391, 392, 393], "c": [1, 4, 146, 258, 259, 260, 261, 262, 264, 265, 270, 271, 274, 275, 390, 391, 393], "scalar": [1, 12, 14, 27, 46, 73, 76, 80, 82, 101, 102, 104, 130, 131, 132, 133, 134, 137, 144, 145, 146, 148, 154, 155, 156, 157, 159, 161, 173, 175, 181, 186, 192, 195, 196, 203, 223, 237, 241, 244, 342, 387, 389, 391], "i": [1, 2, 4, 5, 6, 7, 14, 16, 25, 34, 73, 82, 85, 86, 87, 88, 91, 92, 93, 94, 98, 99, 102, 107, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 130, 131, 137, 142, 143, 146, 147, 149, 154, 158, 159, 163, 166, 167, 169, 181, 182, 184, 185, 186, 187, 190, 191, 194, 195, 196, 199, 202, 203, 204, 209, 213, 215, 220, 221, 226, 227, 230, 233, 237, 238, 239, 240, 241, 245, 246, 250, 256, 258, 259, 260, 261, 262, 263, 264, 265, 267, 268, 269, 270, 271, 272, 273, 274, 275, 281, 282, 288, 290, 291, 293, 294, 296, 297, 298, 299, 300, 301, 303, 307, 309, 310, 311, 315, 316, 321, 322, 323, 329, 330, 332, 337, 338, 341, 342, 344, 349, 353, 359, 363, 366, 368, 369, 370, 375, 377, 378, 383, 386, 387, 388, 389, 390, 391, 392, 393, 394], "sum": [1, 3, 12, 94, 136, 146, 158, 213, 230, 256, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 388, 390], "element": [1, 11, 12, 17, 18, 19, 20, 21, 22, 25, 66, 81, 89, 90, 91, 92, 93, 94, 97, 101, 102, 104, 105, 106, 108, 110, 111, 129, 130, 133, 134, 137, 138, 139, 140, 141, 144, 145, 150, 151, 152, 153, 154, 155, 156, 157, 161, 163, 173, 175, 176, 182, 184, 185, 197, 198, 201, 209, 210, 211, 212, 216, 217, 223, 226, 228, 229, 232, 237, 241, 263, 264, 265, 268, 271, 276, 298, 301, 303, 325, 327, 328, 343, 344, 346, 349, 350, 351, 386, 387], "wise": [1, 11, 12, 17, 18, 19, 20, 21, 22, 81, 89, 90, 101, 102, 104, 105, 106, 108, 110, 129, 130, 133, 134, 137, 144, 145, 150, 151, 152, 153, 154, 155, 156, 157, 161, 173, 175, 176, 197, 201, 209, 210, 211, 212, 216, 217, 223, 228, 229, 264, 265, 276, 298, 325, 327, 328, 343, 344, 346, 349, 350, 351, 386], "numpi": [1, 4, 5, 6, 12, 14, 16, 80, 101, 102, 104, 133, 134, 137, 144, 145, 154, 159, 161, 173, 175, 223, 389, 391, 392], "style": [1, 12, 14, 101, 102, 104, 133, 134, 137, 144, 145, 154, 159, 161, 173, 175, 223], "broadcast": [1, 12, 14, 80, 82, 101, 102, 104, 131, 133, 134, 137, 144, 145, 154, 159, 161, 173, 175, 186, 187, 190, 195, 196, 223, 227, 241, 297], "between": [1, 6, 82, 128, 310, 330, 333, 334, 337, 377, 389, 393], "input": [1, 3, 4, 11, 12, 13, 14, 15, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 98, 99, 101, 102, 104, 105, 106, 108, 109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 132, 133, 134, 136, 137, 138, 139, 140, 141, 143, 144, 145, 146, 147, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 172, 173, 174, 175, 176, 179, 180, 181, 182, 183, 184, 185, 194, 197, 198, 199, 200, 201, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 235, 236, 237, 238, 240, 241, 243, 258, 259, 260, 261, 262, 264, 265, 266, 268, 269, 270, 271, 272, 273, 274, 275, 297, 299, 300, 301, 303, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 324, 329, 330, 332, 333, 334, 335, 337, 338, 340, 342, 353, 356, 386, 387, 388, 391, 392], "upcast": 1, "const": [1, 332], "factor": [1, 147, 311, 331, 376, 379], "streamordevic": 1, "stream": [1, 6, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 44, 45, 48, 49, 50, 51, 52, 53, 54, 55, 56, 59, 60, 61, 62, 63, 65, 67, 68, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 96, 97, 98, 99, 101, 102, 104, 105, 106, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 133, 134, 135, 136, 137, 138, 139, 140, 141, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 190, 191, 192, 194, 195, 196, 197, 198, 199, 200, 201, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 238, 241, 242, 243, 393], "schedul": [1, 169, 361, 375, 376, 377, 378, 379, 381, 393], "itself": [1, 370], "call": [1, 2, 4, 5, 28, 130, 256, 266, 282, 294, 305, 359, 361, 370, 386, 387, 389], "other": [1, 4, 6, 142, 146, 256, 283, 359, 368, 386, 388, 389, 391], "within": [2, 25, 137], "simplest": [1, 256], "wai": [1, 4, 7, 256, 311, 386, 387, 388], "about": [1, 4, 5, 389, 393], "term": [1, 332, 362, 363, 364, 365, 366, 367, 373], "exist": [1, 2, 4, 282, 294], "auto": [1, 7], "ax": [1, 13, 15, 23, 24, 74, 109, 117, 118, 120, 121, 123, 124, 126, 127, 128, 136, 146, 158, 160, 162, 172, 181, 183, 213, 218, 220, 224, 225, 230, 233, 238, 387], "multipli": [1, 35, 184, 185, 263, 307, 311], "earlier": 1, "goal": [], "themselv": [1, 386], "contain": [1, 4, 25, 26, 64, 83, 99, 122, 123, 124, 146, 155, 156, 157, 184, 215, 241, 256, 281, 283, 284, 290, 310, 338, 356, 359, 386, 387], "act": [1, 337], "data": [1, 5, 6, 9, 16, 111, 125, 126, 131, 135, 148, 178, 195, 234, 242, 265, 312, 313, 314, 315, 316, 317, 318, 319, 386, 388, 390], "nor": [1, 132, 237], "rather": [1, 387, 393], "easi": [1, 256], "interfac": 1, "block": [1, 4, 310], "A": [1, 4, 6, 7, 8, 64, 76, 83, 112, 113, 115, 132, 143, 146, 147, 149, 158, 159, 160, 172, 184, 186, 187, 188, 190, 191, 192, 195, 196, 215, 219, 222, 237, 239, 240, 244, 245, 246, 247, 248, 256, 260, 264, 268, 269, 270, 272, 281, 285, 286, 289, 295, 296, 300, 305, 307, 310, 313, 314, 316, 323, 342, 343, 359, 361, 365, 367, 369, 370, 372, 377, 386, 387, 389, 390], "It": [1, 4, 7, 132, 208, 237, 256, 296, 299, 369, 381, 390, 392], "creat": [1, 4, 7, 111, 135, 222, 256, 359, 361, 377, 386, 388, 390], "output": [1, 4, 7, 13, 14, 15, 16, 25, 80, 83, 91, 92, 93, 94, 111, 112, 113, 114, 115, 122, 125, 126, 127, 131, 132, 135, 137, 146, 148, 158, 160, 162, 163, 172, 178, 179, 182, 183, 186, 187, 188, 190, 191, 192, 195, 196, 205, 206, 213, 218, 220, 224, 227, 234, 237, 238, 239, 240, 241, 242, 243, 258, 259, 260, 261, 262, 270, 273, 274, 275, 297, 299, 309, 310, 311, 313, 314, 315, 316, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 353, 356, 386, 387, 388, 389, 390, 391, 392, 393], "given": [1, 13, 15, 25, 35, 80, 82, 84, 91, 92, 93, 94, 97, 99, 107, 109, 116, 117, 118, 119, 120, 121, 125, 126, 127, 131, 146, 158, 160, 162, 168, 172, 177, 183, 190, 192, 200, 208, 213, 215, 220, 224, 231, 232, 234, 235, 236, 238, 248, 258, 259, 263, 274, 275, 281, 297, 330, 332, 338], "set": [1, 4, 5, 7, 83, 100, 103, 112, 114, 168, 169, 207, 208, 222, 267, 272, 273, 280, 282, 289, 290, 291, 294, 295, 299, 303, 309, 330, 342, 353, 359, 363, 370, 383, 387, 389], "further": [1, 7, 387], "class": [1, 4, 5, 8, 9, 10, 27, 248, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 331, 359, 362, 363, 364, 365, 366, 367, 368, 373, 374, 381], "under": [1, 146], "These": [1, 83, 227, 331, 393], "word": [], "bit": [97, 184, 185, 250, 277, 299, 300], "abstract": [], "back": [4, 167, 390], "give": [1, 4, 5, 25, 386], "ourselv": [], "concret": [1, 268, 271, 273, 301, 389, 393], "imag": [262, 264, 265, 311], "public": [1, 256], "explicit": [1, 370, 383, 390], "alpha_": 1, "beta_": 1, "must": [1, 2, 7, 82, 131, 146, 186, 187, 190, 192, 195, 196, 241, 311, 390], "know": [1, 4], "popul": 1, "To": [1, 2, 3, 4, 5, 7, 168, 256, 356, 386, 387, 391], "avoid": [1, 289, 386], "unnecessari": [1, 4], "alloc": [1, 165, 168, 169, 359], "respons": 1, "space": [1, 148, 340], "void": 1, "eval_cpu": 1, "std": [1, 318], "overrid": [1, 103], "eval_gpu": 1, "jacobian": [1, 143, 239, 391], "product": [1, 93, 136, 143, 159, 180, 183, 230, 239, 297, 391], "primal": [1, 143, 239], "tangent": [1, 21, 22, 143, 228, 229, 354], "int": [1, 4, 5, 8, 13, 15, 16, 23, 24, 25, 26, 30, 31, 32, 33, 37, 38, 39, 40, 41, 42, 45, 52, 53, 54, 55, 56, 59, 62, 64, 67, 70, 71, 72, 73, 75, 80, 84, 85, 86, 87, 91, 92, 93, 94, 97, 98, 99, 109, 111, 114, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 131, 132, 135, 142, 146, 148, 158, 160, 162, 164, 165, 166, 168, 169, 172, 174, 178, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 198, 199, 200, 213, 214, 215, 218, 219, 220, 224, 225, 226, 227, 230, 231, 232, 233, 234, 235, 236, 237, 238, 240, 242, 248, 256, 258, 259, 260, 261, 262, 266, 268, 269, 270, 271, 272, 273, 274, 275, 297, 299, 300, 301, 303, 307, 310, 324, 330, 331, 335, 340, 342, 359, 375, 377, 378, 379], "argnum": [1, 132, 237, 387], "cotan": 1, "across": [1, 269], "pair": [1, 181, 284, 303], "repres": [1, 4, 338, 342, 390], "axi": [1, 4, 5, 13, 15, 23, 24, 25, 26, 30, 31, 32, 33, 37, 38, 39, 40, 52, 53, 54, 55, 59, 67, 70, 71, 75, 84, 91, 92, 93, 94, 99, 109, 112, 113, 116, 119, 122, 123, 124, 125, 126, 127, 128, 146, 158, 160, 162, 172, 174, 181, 182, 183, 187, 198, 213, 214, 215, 218, 219, 220, 224, 225, 226, 227, 231, 232, 233, 238, 240, 258, 259, 274, 275, 301, 324, 328, 330, 331, 335, 340, 342, 350, 388], "correspond": [1, 13, 15, 73, 82, 97, 99, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 158, 160, 172, 183, 224, 230, 240, 246, 387], "dimens": [1, 4, 13, 15, 23, 24, 58, 64, 73, 77, 78, 79, 83, 86, 87, 99, 109, 114, 123, 124, 126, 127, 128, 136, 146, 147, 158, 159, 160, 162, 172, 183, 184, 187, 194, 220, 224, 227, 230, 233, 238, 260, 261, 262, 264, 265, 268, 269, 270, 271, 272, 297, 300, 301, 303, 310, 311, 324, 331, 386, 387], "vmap": [1, 387, 389, 391], "print": [1, 2, 3, 4, 5, 7, 245, 246, 247, 256, 383, 386, 387, 388, 389, 390, 391], "ostream": 1, "o": [1, 7, 115, 271], "equival": [1, 28, 61, 74, 102, 130, 226, 267, 296, 298, 299, 302, 304, 306, 308], "check": [1, 7, 76, 142, 167, 284, 387, 388], "bool": [1, 13, 14, 15, 23, 24, 30, 31, 32, 33, 37, 38, 39, 40, 52, 53, 54, 55, 59, 71, 73, 75, 76, 83, 87, 91, 92, 93, 94, 114, 137, 142, 146, 149, 158, 160, 162, 163, 167, 169, 170, 172, 183, 185, 220, 224, 238, 260, 261, 262, 268, 269, 270, 271, 272, 273, 277, 281, 282, 284, 289, 291, 294, 297, 299, 301, 303, 307, 310, 311, 329, 332, 363, 374], "is_equival": 1, "privat": 1, "fall": 1, "eval": [1, 2, 3, 4, 5, 256, 359, 361, 386, 387, 389, 391], "deriv": [1, 387, 389], "base": [1, 114, 146, 151, 153, 303, 310, 359, 361, 367, 381, 383, 386, 388], "abov": [1, 4, 184, 235, 256, 311, 366, 387, 388, 389, 393], "demonstr": 390, "treat": [1, 123, 124, 126, 127, 226, 311, 386], "paramet": [1, 3, 4, 5, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 34, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 97, 98, 99, 101, 102, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 168, 169, 170, 172, 173, 174, 175, 176, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 277, 278, 281, 282, 284, 289, 290, 291, 294, 295, 296, 297, 298, 299, 300, 301, 303, 305, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 324, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 353, 355, 356, 359, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 372, 373, 374, 375, 376, 377, 378, 379, 381, 386, 387, 389], "produc": [1, 83, 297, 356], "through": [1, 221, 310, 368, 386, 387, 390], "construct": [1, 5, 41, 98, 131, 178, 231, 242], "its": [1, 7, 159, 182, 194, 234, 244, 247, 256, 299, 365, 366, 367, 390, 393], "type": [1, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 34, 64, 73, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 97, 98, 99, 101, 102, 104, 105, 106, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 168, 169, 170, 172, 173, 174, 175, 176, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 194, 195, 196, 197, 198, 199, 200, 201, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 245, 256, 289, 310, 312, 313, 314, 315, 316, 317, 318, 319, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 386, 388], "shape": [1, 2, 4, 5, 61, 76, 80, 83, 85, 86, 87, 99, 115, 116, 119, 122, 125, 126, 127, 131, 143, 159, 178, 179, 186, 187, 188, 190, 191, 192, 195, 196, 199, 227, 239, 241, 242, 243, 256, 258, 259, 260, 261, 262, 264, 265, 268, 270, 271, 273, 274, 275, 284, 301, 312, 313, 314, 315, 316, 317, 318, 319, 331, 342, 361, 386, 387, 388, 391, 393], "pass": [1, 4, 5, 61, 74, 180, 181, 237, 244, 245, 246, 256, 282, 294, 295, 296, 299, 305, 386, 389], "re": [5, 7, 356], "now": [1, 4, 7, 299, 386, 390], "promot": 1, "dtype": [1, 4, 10, 16, 27, 34, 35, 73, 111, 128, 131, 135, 142, 146, 147, 148, 178, 188, 190, 191, 192, 195, 196, 234, 242, 250, 289, 311, 312, 313, 314, 315, 316, 317, 318, 319, 329, 331, 338, 375, 376, 377, 378, 379, 386, 387, 388, 390, 391, 392], "promoted_dtyp": 1, "promote_typ": 1, "float32": [1, 10, 16, 111, 115, 135, 142, 146, 147, 148, 178, 188, 190, 191, 195, 196, 234, 242, 250, 311, 312, 313, 314, 315, 316, 317, 318, 319, 329, 331, 338, 375, 376, 377, 378, 379, 386, 387, 388, 389, 390, 391, 392], "non": [1, 7, 163, 292, 301, 343, 359], "point": [1, 3, 4, 7, 130, 185, 250], "out_dtyp": 1, "is_floating_point": 1, "cast": [1, 34, 125, 126, 127, 149, 277, 289, 390], "up": [1, 4, 299, 386], "determin": [1, 99, 190, 250, 288, 392], "x_cast": 1, "astyp": [1, 4, 277, 390], "y_cast": 1, "broadcasted_input": 1, "broadcast_arrai": 1, "out_shap": 1, "0": [1, 3, 4, 5, 7, 8, 16, 35, 41, 42, 45, 62, 67, 75, 84, 85, 86, 87, 98, 99, 111, 115, 128, 132, 146, 147, 168, 181, 186, 191, 196, 198, 200, 215, 219, 220, 234, 235, 236, 237, 238, 240, 245, 256, 258, 259, 260, 261, 262, 263, 264, 265, 267, 269, 270, 272, 274, 275, 298, 302, 303, 307, 308, 309, 310, 312, 313, 314, 315, 316, 317, 318, 319, 320, 322, 323, 325, 326, 329, 331, 333, 334, 338, 341, 342, 344, 345, 346, 347, 352, 353, 356, 359, 362, 363, 365, 366, 367, 368, 370, 373, 374, 375, 376, 377, 378, 379, 383, 386, 387, 388, 389, 390, 391, 392], "unique_ptr": 1, "make_shar": 1, "to_stream": 1, "handl": [1, 256, 386], "resolv": 1, "No": [1, 4], "happen": [1, 4, 112, 310, 361, 386, 389], "alon": [1, 390], "effect": [264, 386, 389], "onli": [1, 4, 6, 7, 76, 85, 86, 87, 146, 184, 190, 256, 281, 282, 284, 289, 291, 294, 295, 296, 359, 386, 387, 392, 393], "execut": [1, 7, 77, 78, 79, 166, 390, 393], "depend": [1, 2, 3, 73, 146, 268, 271, 301, 388, 392, 393], "devic": [1, 6, 7, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 44, 45, 48, 49, 50, 51, 52, 53, 54, 55, 56, 59, 60, 61, 62, 63, 65, 67, 68, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 101, 102, 104, 105, 106, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 133, 134, 135, 136, 137, 138, 139, 140, 141, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 169, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 190, 191, 192, 194, 195, 196, 197, 198, 199, 200, 201, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 238, 241, 242, 243, 248, 393, 394], "specifi": [1, 16, 34, 86, 87, 99, 123, 124, 131, 132, 146, 148, 174, 178, 187, 198, 225, 226, 227, 230, 233, 237, 240, 242, 260, 309, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 353, 387, 393], "memori": [1, 6, 164, 165, 166, 168, 169, 310, 359, 363, 386, 389, 390], "ha": [1, 2, 4, 5, 6, 73, 83, 99, 122, 123, 125, 126, 127, 132, 163, 165, 187, 260, 268, 271, 273, 301, 359, 361, 386, 388, 389, 391, 393], "been": [1, 4, 165, 389], "try": [1, 7], "naiv": [1, 387], "gener": [1, 2, 3, 10, 16, 87, 111, 123, 124, 148, 163, 186, 190, 191, 192, 195, 196, 310, 383, 386, 388, 389, 394], "version": [1, 7, 97, 154, 158, 184, 213, 240, 383, 387, 388], "declar": 1, "member": [1, 256, 287, 292], "method": [1, 4, 8, 9, 10, 27, 248, 256, 288, 359, 362, 363, 364, 365, 366, 367, 368, 370, 373, 374, 381], "each": [1, 64, 97, 107, 114, 142, 159, 163, 181, 184, 185, 187, 198, 205, 206, 215, 231, 233, 240, 241, 264, 265, 266, 268, 269, 271, 301, 303, 310, 329, 331, 383, 386, 389], "find": [1, 3, 7], "pointwis": [], "captur": [1, 2, 83, 170, 171, 256, 386], "templat": 1, "axpby_impl": 1, "typenam": 1, "t": [1, 4, 105, 115, 185, 237, 256, 258, 268, 271, 274, 301, 362, 363, 364, 365, 366, 367, 368, 373, 374, 386, 387, 393], "readi": 1, "fill": [1, 131, 179, 234, 243, 312, 313, 314, 315, 316, 318, 319], "malloc_or_wait": 1, "synchron": [1, 386], "avail": [1, 3, 4, 5, 7, 9, 167, 393], "There": [1, 256, 311, 386], "wait": [1, 4, 169], "here": [1, 4, 386, 387, 389, 392, 393], "request": 1, "pressur": 1, "condit": [1, 241, 393], "set_data": 1, "nbyte": 1, "collect": [1, 246, 385], "pointer": 1, "x_ptr": 1, "y_ptr": 1, "out_ptr": 1, "relev": 1, "static_cast": 1, "size_t": 1, "out_idx": 1, "size": [1, 4, 5, 47, 64, 86, 97, 109, 112, 113, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 131, 135, 142, 146, 165, 169, 184, 185, 187, 199, 215, 218, 256, 258, 259, 261, 262, 266, 270, 274, 275, 299, 311, 363, 389, 390], "map": [1, 5, 35, 149, 246, 266, 277], "linear": [1, 4, 5, 6, 246, 256, 267, 284, 299, 301, 302, 304, 306, 311, 320, 321, 322, 323, 324, 326, 345, 346, 347, 349, 356, 359, 370, 378, 386], "indic": [1, 14, 23, 24, 25, 26, 35, 132, 137, 138, 139, 140, 141, 215, 226, 227, 237, 291, 293, 331, 338, 377, 388], "offset": [1, 4, 42, 99, 112, 114], "x_offset": 1, "elem_to_loc": 1, "stride": [1, 85, 86, 87, 258, 259, 261, 262, 274, 275, 303, 388], "y_offset": 1, "contigu": 1, "regularli": 1, "default": [1, 7, 13, 14, 15, 16, 23, 24, 25, 26, 76, 83, 84, 85, 86, 87, 95, 96, 97, 98, 99, 111, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 132, 135, 137, 146, 147, 148, 149, 158, 160, 162, 163, 168, 169, 172, 178, 182, 183, 184, 185, 186, 187, 188, 190, 191, 192, 194, 195, 196, 198, 199, 200, 207, 208, 214, 215, 218, 219, 220, 222, 224, 230, 232, 233, 234, 235, 236, 237, 238, 240, 242, 250, 258, 259, 260, 261, 262, 268, 270, 271, 273, 274, 275, 277, 282, 284, 289, 291, 294, 297, 298, 299, 301, 303, 307, 308, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 324, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 359, 362, 363, 364, 365, 366, 367, 368, 373, 374, 375, 383, 385, 386, 387, 390, 392, 394], "row": [1, 111, 135, 184, 234], "major": 1, "henc": [1, 184, 386], "doesn": [1, 256], "addit": [1, 4, 12, 112, 113, 115, 149, 260, 269, 272, 297, 300, 359, 387], "abl": [1, 184], "all": [1, 2, 5, 7, 14, 25, 35, 77, 78, 79, 83, 86, 87, 111, 118, 121, 124, 127, 159, 181, 182, 218, 256, 277, 278, 282, 285, 286, 287, 292, 294, 297, 299, 307, 310, 311, 356, 359, 381, 383, 386, 388, 389, 391, 394], "incom": 1, "accordingli": 1, "dispatch": 1, "float16": [1, 10, 149, 250, 277, 389, 390], "bfloat16": [1, 10, 250, 390], "complex64": [1, 250], "throw": [1, 83], "error": [1, 7, 105, 106, 169, 215, 267, 299, 321, 322, 323, 337, 339, 387, 390], "encount": [1, 387], "unexpect": [1, 16], "regist": [1, 5], "op": [1, 180, 282, 389], "assert": 1, "2": [1, 3, 4, 5, 35, 86, 98, 99, 105, 117, 120, 122, 123, 124, 125, 126, 127, 128, 142, 146, 147, 153, 159, 184, 190, 194, 230, 234, 235, 236, 250, 256, 258, 259, 262, 267, 274, 275, 300, 307, 311, 312, 313, 314, 315, 316, 317, 318, 319, 322, 331, 332, 334, 341, 342, 356, 359, 362, 364, 365, 366, 370, 373, 386, 387, 388, 389, 390, 391, 392, 393], "1": [1, 2, 4, 5, 16, 25, 26, 35, 42, 45, 85, 86, 87, 98, 99, 110, 115, 116, 117, 119, 120, 122, 123, 124, 125, 126, 127, 128, 136, 142, 146, 147, 159, 163, 169, 180, 182, 184, 187, 190, 191, 196, 209, 214, 226, 232, 237, 250, 256, 258, 259, 260, 261, 262, 263, 264, 265, 267, 268, 269, 270, 271, 272, 273, 274, 275, 298, 300, 301, 303, 307, 309, 311, 313, 314, 315, 316, 317, 318, 319, 320, 322, 323, 324, 327, 328, 329, 330, 331, 332, 333, 334, 335, 337, 338, 340, 341, 342, 347, 348, 350, 351, 353, 356, 359, 361, 362, 363, 364, 365, 366, 367, 368, 370, 373, 374, 375, 376, 377, 378, 379, 386, 387, 388, 390, 391, 392, 393], "correct": [1, 7, 365, 366, 367, 388, 389], "els": [1, 4, 256, 282, 389], "float16_t": 1, "bfloat16_t": 1, "complex64_t": 1, "runtime_error": 1, "support": [1, 4, 6, 7, 14, 85, 86, 87, 115, 128, 137, 147, 149, 159, 184, 190, 387, 388, 390, 392], "have": [1, 4, 7, 14, 76, 77, 78, 79, 123, 124, 126, 127, 137, 159, 170, 187, 245, 271, 297, 305, 368, 370, 385, 386, 388, 389, 393], "rememb": [], "3": [1, 4, 7, 128, 142, 146, 147, 311, 314, 316, 325, 363, 368, 383, 386, 388, 390, 391], "complic": [], "keep": [1, 13, 15, 23, 24, 158, 160, 162, 172, 183, 220, 224, 238, 256, 281, 387, 389], "mind": [1, 4], "half": [1, 16, 192, 196, 303, 389], "precis": [1, 4, 110, 115, 256, 267, 300, 369, 386], "direct": [1, 4, 279, 368, 393], "fix": [1, 4, 7, 389], "possibli": [4, 159], "due": [], "transpos": [4, 28, 185], "aren": [], "guarante": [], "fit": [1, 184, 393], "requir": [1, 4, 256, 389, 390], "column": [1, 111, 135, 184], "inplac": 1, "expect": [1, 4, 261, 262, 263, 264, 265, 307, 310, 332, 386, 388], "answer": [], "copi": [1, 4, 6, 182, 214, 390], "simpli": [1, 4, 7, 302, 320, 326, 345, 354, 359, 386, 387], "catlas_saxpbi": 1, "axpby_impl_acceler": 1, "first": [1, 2, 3, 4, 5, 7, 99, 128, 132, 155, 157, 159, 182, 194, 225, 230, 237, 245, 256, 259, 269, 275, 311, 330, 338, 363, 365, 366, 367, 370, 386, 387, 390, 393], "mode": [1, 88, 280, 291, 293, 311, 315, 316], "e": [1, 5, 7, 105, 143, 209, 260, 261, 262, 264, 265, 269, 270, 272, 282, 300, 327, 328, 350, 355, 361, 364, 386, 389, 394], "match": [7, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 164, 284, 311, 331, 388, 390], "transposit": [], "data_s": [], "items": [], "flag": [1, 386, 390], "copy_inplac": 1, "copytyp": 1, "n": [1, 4, 27, 85, 86, 87, 111, 116, 118, 119, 121, 122, 125, 127, 135, 190, 220, 234, 238, 258, 259, 260, 261, 262, 264, 265, 268, 271, 274, 275, 301, 311, 337, 342], "incx": 1, "inci": 1, "great": 2, "But": 393, "criteria": 1, "luckili": 389, "alwai": [164, 245, 387], "With": 1, "final": [1, 3, 4, 5, 378], "singl": [1, 5, 107, 143, 149, 163, 181, 239, 259, 275, 386, 388, 392], "row_contigu": 1, "col_contigu": 1, "common": [1, 361, 386, 389], "hit": 1, "mileston": [], "enough": [1, 389], "run": [1, 2, 4, 5, 6, 7, 8, 180, 248, 260, 277, 362, 363, 365, 366, 367, 386, 389, 393, 394], "If": [1, 4, 7, 13, 14, 15, 16, 23, 24, 25, 26, 73, 76, 82, 84, 88, 91, 92, 93, 94, 98, 99, 107, 112, 114, 125, 126, 127, 130, 131, 132, 137, 146, 149, 158, 159, 160, 162, 163, 168, 169, 172, 178, 181, 182, 183, 187, 190, 198, 213, 214, 215, 220, 224, 226, 227, 230, 232, 237, 238, 240, 242, 246, 260, 261, 262, 269, 272, 273, 282, 284, 294, 299, 301, 303, 305, 307, 311, 329, 331, 342, 363, 386, 387, 389, 392, 393, 394], "plan": [1, 386], "stop": [1, 4, 16, 148, 171, 221, 387, 388], "enjoi": 1, "speed": 1, "appl": [1, 4, 6, 7, 393], "silicon": [1, 4, 6, 7, 393], "address": 1, "shade": 1, "languag": 1, "kernel": [1, 85, 86, 87, 258, 259, 274, 275, 386, 388], "written": 1, "help": [1, 4, 386, 393], "resourc": 1, "walkthrough": 1, "pipelin": 1, "specif": [1, 7, 387], "cpp": 1, "algorithm": [311, 368], "launch": [1, 388], "exactli": [1, 4, 284, 387], "mani": [1, 215, 261, 262, 266, 386, 389], "thread": 1, "pick": 1, "updat": [1, 3, 4, 5, 35, 83, 246, 260, 277, 278, 284, 289, 290, 291, 296, 361, 363, 366, 368, 369, 370, 374, 375, 376, 377, 378, 379, 386, 389], "assign": [1, 35, 359], "axpby_gener": 1, "buffer": [1, 164, 390], "constant": [1, 4, 7, 112, 113, 181, 256, 260, 269, 272, 300, 332, 342, 373, 375, 386, 390], "4": [1, 4, 97, 128, 146, 184, 185, 205, 250, 258, 259, 260, 270, 274, 275, 299, 310, 311, 313, 314, 315, 329, 386, 388, 391, 393], "5": [1, 3, 4, 7, 146, 169, 186, 258, 260, 263, 264, 265, 270, 274, 308, 311, 312, 315, 316, 341, 352, 356, 373, 375, 376, 386, 387, 388], "x_stride": 1, "6": [1, 4, 146, 205, 310, 314, 322, 323, 325, 332, 342, 346, 373, 386, 388, 391], "y_stride": 1, "7": [1, 4, 146, 184, 388], "ndim": [1, 128, 146, 311], "8": [1, 4, 7, 146, 184, 250, 259, 270, 275, 310, 330, 362, 363, 364, 365, 366, 367, 373, 386, 388, 391, 393], "uint": 1, "index": [1, 6, 8, 25, 35, 109, 111, 132, 163, 182, 226, 227, 237, 248], "thread_position_in_grid": 1, "convert": [1, 73, 77, 78, 79, 128, 299, 389, 390, 391], "instanti": [1, 5, 389], "uniqu": [1, 383], "host": 1, "name": [1, 149, 184, 185, 203, 204, 205, 206, 256, 269, 281, 284, 286, 388, 392], "identifi": [1, 245, 385], "instantiate_axpbi": 1, "type_nam": 1, "host_nam": 1, "axpby_general_": 1, "compil": [2, 6, 7, 100, 103, 387, 389], "mlx_ext": 1, "metallib": [1, 7], "see": [1, 4, 5, 7, 9, 10, 29, 30, 31, 32, 33, 36, 37, 38, 39, 40, 42, 44, 45, 48, 49, 50, 51, 52, 53, 54, 55, 56, 59, 60, 61, 62, 63, 65, 67, 68, 69, 70, 71, 72, 74, 75, 146, 168, 203, 204, 250, 256, 260, 264, 267, 280, 298, 299, 302, 303, 304, 306, 307, 308, 311, 313, 314, 315, 316, 321, 322, 323, 347, 386, 387, 388, 391, 393], "later": [2, 7], "co": [1, 307, 387], "locat": [1, 295, 296, 393], "share": [6, 97, 184, 185], "register_librari": 1, "potenti": 169, "path": [2, 7, 170, 205, 206, 284], "tri": [], "load": [5, 6, 284], "hasn": [], "alreadi": [1, 2, 4], "static": 7, "object": [2, 9, 27, 46, 73, 83, 142, 205, 240, 245, 246, 250, 264, 310, 385], "why": 4, "packag": [1, 3, 5, 356], "process": [4, 87, 88, 246, 265, 266, 310, 385], "logic": [1, 155, 156, 157], "grid": [1, 163], "shown": 1, "below": [1, 7, 146, 234, 236, 250, 311, 389], "prepar": [1, 4], "carri": 1, "should": [1, 3, 4, 5, 7, 99, 112, 113, 115, 143, 170, 184, 227, 237, 239, 245, 256, 261, 262, 264, 265, 291, 297, 305, 331, 333, 338, 359, 385, 386, 387, 389, 390, 394], "d": [1, 4, 98, 99, 136, 146, 159, 163, 180, 226, 234, 235, 236, 247, 265, 268, 271, 301, 362, 365, 367, 393], "ostringstream": 1, "kname": 1, "axpby_": 1, "general_": 1, "type_to_nam": 1, "make": [1, 2, 4, 5, 7, 159, 177, 208, 256, 375, 376, 378, 379, 386, 389, 391, 393], "sure": [1, 2, 4, 7, 256, 386], "look": [1, 4], "folder": 1, "get_colocated_mtllib_path": 1, "get_kernel": 1, "str": [1, 88, 132, 146, 149, 163, 170, 202, 203, 204, 205, 206, 237, 245, 247, 277, 278, 281, 282, 284, 286, 288, 294, 311, 315, 316, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342], "encod": [1, 114, 303, 307, 310, 331], "compute_encod": 1, "get_command_encod": 1, "setcomputepipelinest": 1, "those": [1, 4, 256], "nelem": 1, "set_array_buff": 1, "setbyt": 1, "sizeof": 1, "threadgroup": 1, "higher": [1, 136, 338, 387], "than": [1, 4, 73, 88, 99, 102, 114, 133, 134, 144, 145, 159, 168, 246, 303, 309, 311, 338, 341, 353, 363, 368, 386, 387, 393], "max": [1, 146, 161, 274, 275, 298, 325, 330, 332, 333, 338, 342, 344, 346, 363, 367, 386, 387, 393], "allow": [1, 142, 256, 296, 359, 381, 388, 391], "tgp_size": 1, "min": [1, 146, 173, 298, 325, 344, 346], "maxtotalthreadsperthreadgroup": 1, "3d": [1, 260, 265, 311], "mtl": 1, "group_dim": 1, "grid_dim": 1, "divid": [1, 35, 130, 184], "among": 1, "dispatchthread": 1, "few": [1, 4, 5, 6, 389, 391], "thing": [1, 4], "note": [1, 4, 7, 14, 83, 85, 86, 115, 123, 124, 137, 146, 164, 184, 187, 256, 300, 311, 390, 392], "befor": [1, 4, 7, 25, 182, 281, 310, 370, 388, 389], "move": [1, 174, 393], "track": [1, 256, 260], "activ": [1, 7, 164, 264, 309, 310, 343, 352, 353, 355, 386], "u": [1, 273, 296, 381, 389], "command": [1, 2, 7], "instead": [1, 7, 256, 296, 307, 387, 389], "end_encod": 1, "end": [99, 167, 184, 259, 268, 271, 275, 309, 334, 341, 347, 352, 353, 378], "until": [1, 389, 391], "limit": [1, 82, 168, 169, 388], "flush": 1, "enqueu": [], "commit": [], "associ": [1, 205, 206, 389], "suggest": [], "deeper": [], "dive": [], "studi": [], "come": [1, 4, 387], "far": 361, "built": [1, 7, 389], "includ": [1, 91, 92, 93, 94, 164, 165, 169, 272, 278, 290, 299, 332, 386, 387, 388, 391, 392, 394], "forward": [1, 237, 386, 389], "diff": 1, "push": 1, "along": [1, 23, 24, 83, 84, 91, 92, 93, 94, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 146, 198, 213, 215, 219, 226, 227, 230, 231, 232, 256, 301, 324], "similarli": [1, 7, 159, 387, 389], "scale_arr": 1, "contribut": 1, "tangent_x": 1, "tangent_i": 1, "revers": [1, 37, 38, 39, 40, 91, 92, 93, 94, 233, 307], "arg": [1, 4, 9, 10, 107, 205, 206], "push_back": 1, "fulli": [1, 6, 386, 390, 393], "overal": 1, "directori": [1, 4, 7], "extens": [1, 149, 170, 288, 392], "h": [1, 85, 86, 146, 259, 260, 262, 264, 265, 268, 271, 275, 301, 387, 389], "mlx_sample_extens": 1, "__init__": [1, 4, 5, 8, 9, 10, 27, 248, 256, 359], "py": [1, 4, 7], "cmakelist": 1, "txt": 1, "setup": [1, 3, 5, 7, 386], "hold": [1, 4, 9, 10, 146, 386], "instal": 1, "pybind11": [], "sinc": [1, 4, 5, 359, 368, 377, 390, 393], "compon": [1, 4], "etc": [1, 184, 256, 311], "pybind11_modul": [], "m": [1, 4, 7, 111, 146, 234, 258, 259, 274, 275, 362, 386], "doc": [1, 5], "sampl": [1, 3, 4, 148, 186, 187, 188, 190, 192, 195, 196, 313, 314, 315, 316, 318, 319, 332, 338, 342, 383, 386], "_a": 1, "pos_onli": [], "kw_onli": 1, "none": [1, 4, 8, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 44, 45, 48, 49, 50, 51, 52, 53, 54, 55, 56, 59, 60, 61, 62, 63, 65, 67, 68, 69, 70, 71, 72, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 171, 172, 173, 174, 175, 176, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 204, 205, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 240, 241, 242, 243, 245, 246, 248, 258, 259, 267, 274, 275, 277, 281, 282, 289, 294, 297, 301, 307, 310, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 363, 381, 388], "r": [1, 4, 147, 237, 264, 268], "pbdoc": [], "most": [1, 187, 256, 372, 386, 387, 388, 389], "complex": [1, 123, 124, 125, 126, 127, 245, 250, 256, 296, 386, 387], "bell": 1, "whistl": 1, "liter": [1, 311, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342], "string": [1, 390, 392], "modul": [1, 4, 5, 244, 299, 305, 310, 356, 372, 385, 386, 389], "ensur": [1, 7, 337], "caster": 1, "find_packag": 1, "config": 1, "add_librari": 1, "sourc": [1, 2, 56, 174, 233], "target_sourc": 1, "cmake_current_list_dir": 1, "header": 1, "target_include_directori": 1, "target_link_librari": 1, "attach": 1, "conveni": [1, 5, 142], "mlx_build_metallib": 1, "target": [1, 237, 329, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 386], "destin": [1, 56, 174], "automat": [1, 6, 149, 391, 392, 393], "practic": [1, 386], "mlx_build_met": [1, 7], "mlx_ext_metallib": 1, "titl": 1, "include_dir": 1, "project_source_dir": 1, "mlx_include_dir": 1, "output_directori": 1, "cmake_library_output_directori": 1, "add_depend": 1, "endif": 1, "pybind11_add_modul": [], "build_shared_lib": 1, "target_link_opt": 1, "wl": 1, "rpath": 1, "loader_path": 1, "onc": [1, 386], "describ": [1, 389], "util": [1, 4, 6, 7, 205, 256], "__name__": [1, 4], "__main__": [1, 4], "descript": [1, 4, 250], "ext_modul": 1, "cmakeextens": 1, "cmdclass": 1, "build_ext": 1, "cmakebuild": 1, "package_dir": [], "package_data": 1, "dylib": 1, "zip_saf": 1, "fals": [1, 4, 13, 14, 15, 23, 24, 30, 31, 32, 33, 37, 38, 39, 40, 52, 53, 54, 55, 59, 71, 75, 76, 83, 87, 91, 92, 93, 94, 137, 142, 146, 149, 158, 160, 162, 163, 169, 172, 183, 220, 224, 238, 241, 245, 246, 250, 269, 270, 272, 273, 282, 284, 294, 297, 299, 303, 307, 310, 311, 329, 332, 363, 374, 390], "python_requir": 1, "even": [1, 4, 83, 386, 389, 390], "though": [1, 4, 386, 389, 390], "j8": 1, "libmlx_ext": 1, "cpython": 1, "3x": 1, "darwin": 1, "pip": [1, 7], "after": [1, 4, 5, 25, 128, 130, 182, 184, 260, 269, 272, 277, 278, 282, 284, 291, 294, 295, 296, 297, 310, 341, 386, 393], "plai": [1, 4], "ones": [1, 4, 179, 205, 234, 295, 296, 299, 388], "b": [1, 2, 4, 12, 14, 76, 101, 102, 104, 130, 133, 134, 136, 137, 144, 145, 146, 154, 155, 157, 159, 161, 173, 175, 180, 184, 223, 230, 237, 273, 301, 311, 324, 387, 388, 389, 390, 391, 392, 393], "f": [1, 2, 3, 5, 146, 256, 271, 366, 386, 390], "item": [1, 3, 4, 5, 246, 389, 390, 391], "true": [1, 3, 4, 14, 37, 38, 39, 40, 76, 83, 91, 92, 93, 94, 114, 137, 142, 146, 149, 163, 169, 185, 213, 241, 245, 246, 250, 256, 260, 261, 262, 268, 269, 270, 271, 272, 273, 281, 282, 284, 291, 294, 299, 301, 303, 307, 310, 311, 329, 337, 363], "quick": [1, 6], "benchmark": [1, 386], "compar": [1, 76, 386], "time": [1, 4, 7, 169, 231, 256, 258, 259, 268, 271, 274, 275, 301, 386, 387, 389, 393], "set_default_devic": 1, "256": [1, 5], "512": [1, 2, 4, 310, 393], "random": [1, 2, 3, 4, 5, 6, 258, 259, 260, 270, 274, 275, 284, 291, 386, 387, 393, 394], "normal": [1, 3, 4, 112, 113, 190, 195, 256, 258, 259, 260, 269, 270, 272, 274, 275, 300, 310, 313, 315, 390, 393], "bench": 1, "warm": [1, 386], "rang": [1, 2, 3, 4, 5, 7, 16, 128, 148, 314, 316, 322, 323, 361, 375, 376, 377, 378, 379, 383, 386, 387, 389, 393], "100": [1, 3, 4, 378, 386, 387, 389, 393], "5000": 1, "simple_tim": 1, "custom_tim": 1, "3f": [1, 5, 386], "custom": [1, 310], "114": 1, "109": 1, "modest": 1, "improv": [1, 2, 4, 362, 363, 364, 365, 366, 367, 373, 386], "awai": [1, 4], "good": [1, 7, 386, 393], "nn": [1, 4, 5, 205, 246, 256, 356, 359, 361, 370, 372, 386, 389], "grad": [1, 3, 5, 237, 361, 369, 386, 387, 388, 389, 391], "full": [1, 5, 61, 74, 88, 213, 295, 296, 332, 386, 389], "profil": 2, "kei": [2, 4, 115, 186, 187, 188, 190, 191, 192, 194, 195, 196, 245, 246, 281, 282, 294, 297, 370, 383, 385, 387], "build": [2, 4, 6, 315, 359, 386], "mlx_metal_debug": [2, 7], "option": [2, 4, 13, 15, 16, 23, 24, 25, 26, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 44, 45, 48, 49, 50, 51, 52, 53, 54, 55, 56, 59, 60, 61, 62, 63, 65, 67, 68, 69, 70, 71, 72, 74, 75, 77, 78, 79, 83, 84, 85, 86, 87, 88, 91, 92, 93, 94, 97, 98, 99, 111, 112, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 131, 132, 135, 140, 141, 146, 147, 148, 149, 158, 160, 162, 163, 169, 172, 178, 181, 182, 183, 184, 185, 186, 187, 188, 190, 191, 192, 194, 195, 196, 198, 199, 213, 214, 215, 218, 219, 220, 224, 226, 230, 232, 233, 234, 235, 236, 237, 238, 240, 242, 245, 246, 258, 259, 260, 261, 262, 268, 271, 273, 274, 275, 277, 281, 282, 284, 289, 294, 297, 299, 301, 303, 307, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 362, 363, 364, 365, 366, 367, 368, 370, 373, 374, 375, 383, 386, 392, 394], "debug": 2, "record": [2, 166, 389], "dure": [2, 83, 263, 264, 265, 311, 390], "inspect": [2, 386, 391], "label": [2, 3, 331, 338], "queue": 2, "readabl": 2, "start_captur": 2, "initi": [2, 3, 4, 256, 260, 269, 270, 272, 273, 298, 300, 312, 313, 314, 315, 316, 317, 318, 319, 359, 370, 375, 376, 378, 379, 386, 389], "gpu": [2, 6, 386, 388, 393], "main": [6, 99, 111, 246, 256], "jane": [], "develop": [6, 7], "gputrac": [2, 170], "arang": [146, 250, 311, 388, 390], "10": [2, 4, 5, 151, 200, 205, 246, 256, 284, 356, 377, 379, 386, 388], "20": 146, "30": 363, "40": [], "stop_captur": 2, "replai": 2, "trace": [2, 386], "view": [2, 390], "overview": 2, "oper": [2, 4, 6, 8, 34, 77, 78, 79, 87, 115, 213, 221, 227, 248, 256, 310, 368, 386, 387, 388, 389, 390, 391, 393, 394], "checkout": [2, 386], "document": [2, 6, 61, 74, 203, 204, 250, 386, 387, 388], "inform": [2, 4, 5, 7, 203, 204, 250, 256, 260, 267, 297, 387, 393], "skip": 2, "save": [2, 4, 6, 149, 170, 184, 203, 204, 205, 206, 288, 389], "project": [2, 4, 297], "us": [2, 3, 4, 5, 6, 7, 16, 35, 97, 100, 102, 114, 128, 146, 147, 159, 164, 165, 166, 168, 184, 185, 198, 199, 245, 250, 256, 259, 264, 266, 267, 268, 271, 273, 275, 277, 281, 288, 295, 297, 299, 301, 303, 307, 310, 311, 315, 316, 322, 323, 330, 356, 359, 361, 362, 363, 365, 366, 367, 368, 369, 370, 383, 385, 386, 387, 388, 391, 393], "cmake": [2, 7], "mkdir": [2, 7], "cd": [2, 7], "dmlx_metal_debug": 2, "ON": [2, 7], "g": [2, 7, 146, 184, 271, 355, 373, 374, 389, 394], "xcodeproj": 2, "select": [2, 7, 232, 241, 277, 281, 289], "metal_captur": 2, "exampl": [2, 3, 4, 5, 16, 35, 128, 146, 147, 222, 226, 256, 258, 259, 260, 270, 274, 275, 282, 284, 291, 294, 311, 312, 313, 314, 315, 316, 317, 318, 319, 329, 331, 338, 356, 361, 370, 375, 376, 377, 378, 379, 383, 387, 388, 389, 390, 391, 392], "schema": 2, "implement": [3, 5, 114, 115, 146, 266, 281, 297, 303, 305, 307, 309, 310, 311, 353, 362, 363, 364, 365, 367, 368, 369, 381, 386, 387, 390], "basic": [3, 200, 387], "model": [3, 5, 6, 205, 244, 246, 256, 277, 280, 282, 284, 288, 291, 293, 294, 295, 297, 310, 356, 359, 361, 369, 370, 372, 386, 389], "problem": [3, 5, 256], "metadata": [3, 149, 203, 204], "num_featur": [3, 260], "num_exampl": 3, "1_000": 3, "num_it": 3, "10_000": 3, "iter": [3, 5, 246, 383, 386, 389], "sgd": [3, 5, 361, 368, 370, 375, 376, 379, 386], "lr": [3, 368], "01": [3, 326, 366], "rate": [3, 362, 363, 364, 365, 366, 367, 368, 373, 374], "ll": [3, 5, 334, 386, 387], "synthet": 3, "dataset": [3, 389], "matrix": [3, 41, 97, 98, 111, 135, 146, 147, 159, 163, 184, 185, 190, 299, 317, 356], "ground": [3, 4, 331, 341], "truth": [3, 331, 341], "w_star": 3, "valu": [3, 4, 11, 14, 16, 23, 24, 46, 73, 76, 82, 111, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 131, 137, 146, 148, 181, 186, 187, 188, 190, 191, 192, 195, 196, 203, 226, 227, 237, 240, 244, 245, 246, 250, 259, 263, 264, 265, 270, 273, 275, 281, 297, 298, 308, 309, 310, 312, 329, 330, 331, 332, 333, 334, 336, 337, 338, 339, 340, 341, 353, 359, 363, 366, 375, 376, 378, 379, 387], "gaussian": [3, 267, 321, 322, 323, 332], "nois": 3, "noisi": 3, "ep": [3, 112, 113, 260, 269, 270, 272, 300, 330, 332, 342, 362, 363, 364, 365, 366, 367, 373], "1e": [3, 5, 14, 137, 260, 269, 270, 272, 300, 330, 332, 342, 362, 363, 364, 365, 366, 367, 370, 373, 375, 376, 377, 378, 379], "weight": [3, 85, 86, 87, 112, 113, 246, 256, 284, 288, 299, 329, 331, 359, 363, 366, 368, 370, 374, 387, 389], "squar": [3, 4, 113, 135, 201, 216, 237, 246, 256, 300, 339, 341, 362, 363, 365, 366, 367, 387, 390], "loss": [3, 5, 237, 256, 361, 386, 387, 389], "loss_fn": [3, 5, 361, 386, 387], "w": [3, 86, 97, 184, 185, 237, 259, 260, 262, 264, 265, 273, 275, 374, 387], "mean": [3, 4, 5, 113, 190, 191, 237, 256, 260, 269, 282, 300, 318, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 386, 387, 390], "grad_fn": [3, 386, 387], "randomli": [3, 4, 263, 264, 265], "Then": [3, 7], "repeatedli": 3, "_": [2, 3, 4, 256, 375, 376, 377, 378, 379, 383, 386, 389, 393], "verifi": [3, 7], "close": [3, 6, 7, 14, 137], "error_norm": 3, "5f": 3, "someth": [3, 4, 388], "00005": 3, "00364": 3, "complet": [3, 4, 7, 169, 295, 296, 387, 393], "logist": [3, 209, 322, 323, 349], "github": [3, 5, 7, 386], "repo": [3, 5, 7, 386], "enabl": [2, 4, 7, 83, 103, 374], "larg": [4, 256, 297, 337, 386, 389], "ish": 4, "transform": [4, 6, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 244, 256, 260, 269, 272, 273, 281, 282, 294, 299, 303, 388], "compromis": 4, "eas": 4, "llama": 4, "famili": 4, "less": [4, 25, 145, 182, 303, 341], "200": [4, 377], "line": [4, 389, 390], "python": [2, 4, 46, 64, 73, 107, 245, 246, 247, 359, 369, 370, 372, 385, 387, 390], "neural": [4, 6, 266, 313, 314, 343, 356, 359, 373], "network": [4, 6, 260, 264, 266, 313, 314, 356, 359, 373], "concis": 4, "architectur": [4, 7, 256, 296, 393], "notabl": [4, 6], "rope": [4, 256], "posit": [4, 25, 99, 114, 128, 132, 141, 174, 182, 190, 237, 246, 256, 261, 262, 297, 303, 307, 332, 342], "cach": [4, 164, 165, 168, 386], "concaten": 4, "llamaattent": 4, "self": [4, 5, 8, 27, 28, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 44, 45, 46, 48, 49, 50, 51, 52, 53, 54, 55, 56, 59, 60, 61, 62, 63, 65, 67, 68, 69, 70, 71, 72, 73, 74, 75, 248, 256, 343, 359], "dim": [4, 114, 115, 266, 269, 270, 272, 297, 300, 303, 307, 310], "num_head": [4, 297, 310], "super": [4, 5, 256, 359], "tradit": [4, 114, 264, 265, 303], "query_proj": 4, "bia": [4, 97, 112, 184, 185, 246, 256, 261, 262, 268, 271, 272, 273, 282, 284, 294, 297, 299, 301, 365, 366, 367, 370, 387], "key_proj": 4, "value_proj": 4, "out_proj": [4, 359], "__call__": [4, 5, 256, 359], "queri": [4, 115, 297], "mask": [4, 115, 291, 297, 388], "extract": [4, 41, 98, 99, 256, 281, 359], "l": [4, 5, 256, 258, 260, 261, 268, 271, 274, 301, 341], "reshap": [4, 146, 311, 388], "combin": 4, "key_cach": 4, "value_cach": 4, "sqrt": [4, 105, 115, 260, 269, 270, 272, 273, 300, 307, 313, 314, 315, 316, 362, 364, 365, 366, 373, 386], "score": [4, 115, 338], "softmax": [4, 115, 256, 328, 331], "values_hat": 4, "rm": [4, 7, 113, 363], "swiglu": 4, "rmsnorm": [4, 256], "llamaencoderlay": 4, "mlp_dim": [4, 310], "norm1": 4, "norm2": 4, "linear1": 4, "linear2": 4, "linear3": 4, "sigmoid": [4, 256, 306, 322, 323, 327, 349], "instanc": [4, 35, 184, 247, 256, 270, 277, 278, 279, 282, 284, 285, 286, 291, 294, 295, 296, 305, 359, 390], "embed": [4, 256, 303, 307, 330], "emb": [4, 266, 307], "token": [4, 266], "num_lay": [4, 5, 361], "vocab_s": 4, "norm": [4, 113, 269, 342, 367, 368], "multiheadattent": [4, 256], "create_additive_causal_mask": 4, "list": [4, 9, 13, 15, 27, 67, 73, 77, 78, 79, 80, 83, 84, 87, 107, 117, 118, 120, 121, 123, 124, 126, 127, 131, 132, 143, 146, 158, 160, 162, 163, 172, 178, 181, 183, 186, 187, 188, 190, 191, 192, 195, 196, 203, 213, 215, 219, 220, 224, 230, 231, 233, 237, 238, 239, 242, 245, 247, 256, 282, 284, 285, 286, 287, 292, 294, 295, 296, 359, 365, 366, 367, 368, 377, 385, 386, 387, 389], "still": [4, 7, 146, 386, 389], "consid": [4, 14, 76, 137, 245, 246, 269, 385], "train": [4, 5, 256, 260, 263, 264, 265, 280, 282, 294, 313, 314], "ignor": [4, 35, 82, 83, 107, 363], "whatsoev": 4, "rest": [4, 114, 246, 303], "subsect": 4, "prompt": 4, "autoregress": 4, "yield": [4, 5, 383], "temp": 4, "causal": 4, "append": [4, 159, 386, 389], "store": 4, "per": [4, 5, 97, 184, 185, 260, 269, 270, 272, 300, 381, 386, 389], "care": [4, 389], "last": [4, 26, 73, 112, 113, 118, 121, 123, 124, 126, 127, 128, 136, 147, 159, 187, 214, 230, 261, 262, 264, 265, 269, 311, 390], "logit": [4, 187, 329, 331, 386], "next": [1, 4, 5, 168], "categor": 4, "lazili": [4, 256], "noth": [4, 256, 389], "yet": [4, 146, 256, 359, 370, 387, 388, 389, 391], "forc": [4, 5, 256, 391], "choos": [4, 114, 303], "pars": 4, "feed": 4, "loop": [4, 5, 386, 387, 389], "unsqueez": 4, "sequenc": [4, 13, 15, 30, 31, 52, 53, 54, 55, 59, 67, 70, 71, 75, 80, 87, 109, 117, 118, 120, 121, 123, 124, 126, 127, 131, 158, 160, 162, 172, 178, 183, 186, 187, 188, 190, 191, 192, 195, 196, 199, 213, 215, 218, 220, 224, 230, 231, 233, 238, 242, 260, 261, 268, 271, 301, 310, 383, 393], "length": [4, 218, 260, 261, 268, 271, 301, 377], "len": [4, 118, 121, 124, 127, 377], "overwrit": 4, "discard": [4, 245], "old": 4, "moment": [4, 87, 363, 365, 366, 367], "anymor": 4, "everyth": 4, "small": [4, 110, 112, 113, 260, 269, 272, 300, 332, 337, 342, 386, 393], "12": [4, 377], "8192": 4, "1024": 4, "actual": [4, 16, 284, 359, 389], "materi": [4, 6], "could": [4, 256], "20_000": 4, "machin": [4, 6, 7, 373], "8gb": 4, "ram": 4, "32": [4, 5, 184, 185, 250, 259, 275, 300, 386], "44": 4, "doubl": 4, "bracket": 4, "becaus": [4, 164, 256, 389], "batch": [4, 159, 190, 260, 261, 262, 264, 265, 268, 271, 297, 301, 311, 389], "zip": [4, 5], "haven": 4, "anyth": [4, 237, 389], "result": [4, 16, 35, 73, 83, 97, 112, 113, 136, 146, 149, 159, 180, 185, 190, 198, 200, 219, 230, 231, 241, 246, 307, 386, 387, 390], "similar": [4, 142, 246, 295, 296, 297, 330, 390, 392], "runtim": [4, 386], "section": [4, 7, 215, 342, 386, 387], "access": [4, 46, 256, 359, 370, 389, 393], "origin": [4, 99, 260, 290, 313, 314, 315, 316, 362, 363, 364, 365, 367, 368, 390], "sentencepiec": 4, "pytorch": [4, 6, 269, 387], "compat": [4, 187, 190, 392], "npz": [4, 149, 205, 206, 284, 288, 392], "file": [4, 7, 149, 202, 203, 204, 205, 206, 284, 288, 387, 392], "directli": [1, 4], "argpars": 4, "itertool": [4, 246], "starmap": [4, 246], "np": [4, 5, 390, 391], "torch": [4, 390], "map_torch_to_mlx": 4, "tok_embed": 4, "elif": 4, "replac": [4, 295, 296, 310, 341], "attention_norm": 4, "ffn_norm": 4, "wq": 4, "wk": 4, "wv": 4, "wo": 4, "w1": 4, "w2": 4, "w3": 4, "ffn": 4, "separ": [4, 61, 74, 269, 338], "submodul": [4, 5, 256, 278, 282, 283, 294, 296], "feed_forward": 4, "parser": 4, "argumentpars": 4, "add_argu": 4, "torch_weight": 4, "output_fil": 4, "parse_arg": 4, "state": [4, 5, 256, 268, 271, 301, 361, 370, 383, 386], "savez": [4, 288, 392], "k": [4, 41, 98, 111, 115, 232, 234, 235, 236, 258, 273, 274, 282], "v": [4, 88, 115, 256, 282, 390], "left": [4, 114, 146, 184, 258, 259, 267, 274, 275, 303, 311, 322, 323, 332, 334, 342], "disk": 4, "text": [4, 258, 259, 268, 271, 274, 275, 276, 301, 309, 313, 314, 315, 316, 325, 332, 333, 334, 337, 338, 341, 343, 344, 347, 348, 352, 353, 363, 368], "format": [4, 149, 202, 203, 204, 205, 206, 390], "dictionari": [4, 83, 149, 203, 204, 245, 256, 281, 290, 295, 296, 371, 385, 392], "represent": [4, 184, 245, 247], "tree_unflatten": 4, "helper": [4, 386], "weight_fil": 4, "incur": 4, "sever": [4, 85, 86, 87, 205, 206, 386, 392], "futur": [4, 299, 388, 389], "pth": 4, "current": [4, 6, 7, 85, 86, 87, 165, 184, 256, 363, 389], "around": 4, "m1": [4, 386, 387, 393], "ultra": 4, "7b": 4, "me": 4, "ishmael": 4, "year": 4, "ago": 4, "never": [4, 389], "long": 4, "info": [4, 7], "247": 4, "press": [4, 146], "enter": 4, "littl": 4, "monei": 4, "my": [4, 7], "purs": 4, "greater": [4, 25, 110, 134, 182, 309, 353], "consequ": 4, "walk": 4, "down": 4, "gower": 4, "street": 4, "afternoon": 4, "heavi": 4, "rain": 4, "saw": [4, 387], "off": [4, 7, 389], "man": 4, "rag": 4, "who": 4, "sat": 4, "upon": [4, 246], "hi": [4, 271], "bundl": 4, "hard": 4, "wet": 4, "he": [4, 315, 316], "were": [4, 393], "cry": 4, "watch": [4, 386], "him": 4, "observ": 4, "numer": [4, 112, 113, 146, 154, 158, 213, 260, 269, 270, 272, 300, 330, 332, 342, 362, 363, 364, 365, 366, 367, 373, 386, 389], "crowd": 4, "wa": [4, 170, 389], "hurri": 4, "437": 4, "330": 4, "second": [4, 99, 155, 157, 159, 225, 237, 259, 275, 330, 338, 363, 365, 366, 367, 387, 393], "spent": 4, "amount": [4, 166, 258, 274], "39": 4, "By": [4, 289, 387, 390], "bigger": [4, 363], "remain": [4, 237, 263, 264, 265], "almost": 4, "nobodi": 4, "took": 4, "least": [4, 77, 78, 79, 82, 147, 184], "notic": [4, 387, 392], "distanc": [4, 342], "had": 4, "doubt": 4, "minut": 4, "straight": 4, "slowli": 4, "rais": [4, 146, 169, 215, 284], "ey": 4, "speak": [4, 146], "resum": 4, "postur": 4, "stood": 4, "feel": 4, "pain": 4, "heart": 4, "smile": 4, "face": 4, "am": 4, "someon": 4, "three": [4, 79, 311], "quarter": 4, "hour": 4, "made": 4, "immedi": [4, 277], "repli": 4, "again": [4, 7, 256, 386], "hand": [4, 387, 389], "did": 4, "accustom": 4, "thu": [4, 256], "question": [4, 389], "reason": [4, 388], "tell": [4, 386, 390], "understand": [4, 313, 314], "579": 4, "690": 4, "num": [4, 148, 194], "500": [4, 393], "628": 4, "went": 4, "nervou": 4, "trembl": 4, "told": 4, "And": [4, 311], "perhap": [1, 4], "surpris": 4, "matter": [4, 256], "shall": 4, "anyhow": 4, "friend": 4, "ye": 4, "slight": [4, 389], "kind": 4, "longer": [4, 88, 387], "soon": 4, "unless": [4, 14, 137, 146, 359], "unlik": [4, 14, 137, 264, 265, 290], "strang": 4, "amus": 4, "That": 4, "secret": 4, "disappoint": 4, "mine": 4, "cannot": [4, 82, 388, 390], "happi": 4, "ask": 4, "shop": 4, "bui": 4, "food": 4, "633": 4, "21": [4, 379], "475": 4, "su": 4, "j": [4, 7, 146, 264, 364, 365, 367], "lu": 4, "pan": 4, "murtadha": 4, "wen": 4, "liu": 4, "2021": 4, "roform": [4, 303], "enhanc": [4, 303, 389], "rotari": [4, 114, 303], "arxiv": [4, 269, 270, 272, 276, 300, 323, 343, 362, 368], "preprint": [4, 362, 368], "2104": 4, "09864": 4, "zhang": 4, "sennrich": 4, "2019": [4, 366], "root": [4, 113, 201, 216, 300], "advanc": [4, 386], "system": [4, 7, 164, 165], "shazeer": 4, "2020": 4, "glu": [4, 256], "variant": [4, 341, 367], "2002": 4, "05202": 4, "classifi": 5, "mnist": 5, "As": [5, 35, 226, 256, 386], "mlp": [5, 256, 310, 361], "inherit": [5, 385], "standard": [5, 46, 73, 159, 188, 191, 220, 310, 313, 315, 318, 391], "idiom": [5, 386], "input_dim": [5, 256, 273, 299], "hidden_dim": [5, 359, 361], "output_dim": [5, 256, 273, 299], "layer_s": 5, "idim": 5, "odim": 5, "maximum": [5, 23, 35, 82, 91, 166, 169, 256, 302, 307, 322, 323, 326, 345, 359, 389], "cross": [5, 87, 329, 331], "entropi": [5, 329, 331], "sub": [5, 99, 194], "commonli": [5, 295, 356, 386], "cross_entropi": [5, 256], "accuraci": 5, "valid": [5, 88, 128, 240, 245, 282, 294, 385], "eval_fn": 5, "argmax": 5, "loader": 5, "num_class": [5, 361], "batch_siz": [5, 361], "num_epoch": [5, 361], "learning_r": [5, 361, 362, 363, 364, 365, 366, 367, 368, 370, 373, 374, 375, 376, 377, 378, 379, 386], "train_imag": [5, 361], "train_label": [5, 361], "test_imag": 5, "test_label": 5, "shuffl": 5, "minibatch": 5, "batch_iter": [5, 361], "perm": 5, "permut": 5, "id": [5, 7], "put": [5, 386], "trainabl": [5, 244, 256, 359], "loss_and_grad_fn": [5, 361, 386, 387], "value_and_grad": [5, 256, 295, 359, 361, 372, 386, 387, 390, 391], "epoch": 5, "test": [5, 7], "confus": 5, "decent": 5, "95": 5, "brought": 6, "research": 6, "except": [6, 111, 122, 123, 125, 126, 127, 269, 284, 388, 390], "featur": [6, 85, 86, 87, 114, 260, 268, 269, 270, 271, 272, 273, 299, 300, 301, 303, 310, 311, 386, 389], "differ": [6, 142, 223, 341, 387], "lazi": [6, 359, 391], "multi": [6, 115, 261, 262, 388, 390], "cpu": [6, 147, 386, 393], "inspir": 6, "jax": [6, 383], "arrayfir": 6, "unifi": 6, "live": [6, 393], "guid": [1, 6], "convers": 6, "regress": [6, 337], "layer": [6, 112, 256, 258, 259, 264, 265, 268, 269, 271, 272, 273, 274, 275, 291, 296, 299, 301, 305, 310, 355, 359], "perceptron": 6, "llm": 6, "infer": [6, 131, 149], "fft": 6, "algebra": 6, "tree": [6, 83, 107, 132, 237, 240, 245, 246, 247, 369, 370, 372, 381, 387], "debugg": 6, "pypi": 7, "meet": 7, "seri": 7, "chip": 7, "nativ": 7, "maco": 7, "13": 7, "recommend": [7, 169, 368], "14": 7, "sonoma": 7, "conda": 7, "forg": 7, "distribut": [7, 186, 187, 188, 190, 191, 195, 196, 273, 313, 314, 315, 316, 318, 319, 332, 335, 340, 342, 356], "probabl": [7, 192, 263, 264, 265, 299, 329, 331, 335, 393], "platform": 7, "processor": 7, "arm": 7, "i386": 7, "switch": 7, "17": 7, "clang": 7, "24": 7, "xcode": 7, "15": [7, 146, 386], "sdk": 7, "environ": [7, 100, 103], "via": [7, 369, 372, 389, 390], "rosetta": 7, "unam": 7, "p": [7, 186, 256, 263, 264, 265, 342, 365, 367], "clone": 7, "git": 7, "com": 7, "ml": 7, "explor": 7, "nanobind": [1, 7, 310], "http": [7, 269, 270, 272, 276, 300, 323, 343], "wjakob": 7, "env": 7, "cmake_build_parallel_level": 7, "edit": [7, 296], "unittest": 7, "discov": 7, "stub": 7, "dev": [1, 7], "generate_stub": 7, "either": [7, 12, 61, 73, 74, 82, 101, 102, 104, 130, 133, 134, 144, 145, 146, 154, 159, 161, 173, 175, 223, 237, 259, 275, 305, 311, 315, 316], "libmlx": 7, "preprocessor": 7, "metal_path": 7, "mlx_build_test": 7, "mlx_build_exampl": 7, "mlx_build_benchmark": 7, "mlx_build_python_bind": 7, "multipl": [7, 112, 113, 159, 175, 184, 185, 297, 307, 376, 377, 379, 386, 389, 392], "wish": 7, "variabl": [7, 83, 100, 103, 132, 143, 237, 239, 240], "export": 7, "developer_dir": 7, "app": 7, "content": [7, 281, 386], "xcrun": 7, "macosx": 7, "show": [7, 250, 386], "unabl": 7, "tool": 7, "sudo": 7, "ouptut": 7, "finder": 7, "iterm": 7, "termin": 7, "click": 7, "uncheck": 7, "window": [7, 258, 259, 274, 275], "restart": 7, "grep": 7, "cmake_host_system_processor": 7, "arm64": 7, "x86_64": 7, "wipe": 7, "cahc": 7, "rf": 7, "devicetyp": 8, "attribut": [8, 9, 27, 248, 290, 359, 381], "kwarg": [9, 10, 205, 206, 394], "categori": [10, 250], "bool_": [10, 250], "integ": [10, 130, 142, 146, 181, 184, 185, 186, 192, 215, 230, 240, 250, 266, 289, 377, 388], "unsignedinteg": 10, "uint8": [10, 250], "uint16": [10, 250], "uint32": [10, 23, 24, 25, 26, 187, 250], "uint64": [10, 250], "signedinteg": [10, 142], "int8": [10, 250], "int32": [10, 16, 35, 128, 142, 146, 192, 250, 311, 388, 391], "int64": [10, 250], "inexact": [10, 142], "complexflo": 10, "complex128": 10, "issubdtyp": [10, 250], "absolut": [11, 14, 137, 322, 323, 341], "semant": [12, 80, 101, 102, 104, 133, 134, 144, 145, 154, 159, 161, 173, 175, 223, 393], "keepdim": [13, 15, 23, 24, 30, 31, 32, 33, 52, 53, 54, 55, 59, 71, 75, 146, 158, 160, 162, 172, 183, 213, 220, 224, 238], "reduct": [13, 15, 158, 160, 172, 183, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342], "reduc": [13, 15, 23, 24, 158, 160, 162, 172, 183, 220, 224, 238, 260, 310, 337], "unspecifi": [13, 15, 16, 23, 24, 25, 26, 84, 91, 92, 93, 94, 131, 158, 160, 162, 172, 178, 182, 183, 198, 213, 214, 220, 224, 226, 232, 238, 242, 394], "entir": [13, 15, 23, 24, 158, 160, 162, 172, 183, 220, 224, 238, 264, 265], "singleton": [13, 15, 23, 24, 158, 159, 160, 162, 172, 183, 220, 224, 238], "rtol": [14, 137], "05": [14, 137, 260, 269, 270, 272, 300], "atol": [14, 137], "08": [14, 137, 330, 364, 365, 366, 367, 373], "equal_nan": [14, 76, 137], "approxim": [14, 267, 321, 322, 323], "comparison": [14, 104, 133, 134, 144, 145], "infinit": [14, 137], "equal": [14, 25, 76, 111, 134, 137, 145, 182, 192, 215, 270, 273], "sign": [14, 137, 250, 368], "nan": [14, 76, 137, 139], "ab": [14, 137, 146, 237, 269, 270, 272, 276, 300, 323, 343, 386], "array_equ": [14, 137], "rel": [14, 137, 363, 386], "toler": [14, 137], "boolean": [14, 76, 137, 138, 139, 140, 141, 155, 156, 157, 250, 293, 388], "interv": [16, 148, 192, 196], "increment": 16, "otherwis": [16, 87, 169, 245, 246, 282, 284, 294, 309, 310, 311, 329, 334, 341, 352, 353, 389, 390], "convent": [16, 88, 311, 366], "lead": [16, 386], "fraction": 16, "integr": [16, 226, 389], "invers": [17, 18, 19, 20, 21, 22, 106, 119, 120, 121, 122, 123, 124], "cosin": [17, 18, 89, 90, 330, 375, 377, 387], "hyperbol": [18, 20, 22, 90, 212, 229, 354], "sine": [19, 20, 211, 212, 387], "minimum": [24, 35, 82, 92, 307, 330, 375], "kth": [25, 182], "partit": 25, "order": [25, 87, 146, 182, 184, 232, 256, 269, 295, 305, 370, 386, 387], "undefin": [25, 182, 190, 388], "sort": [25, 26, 182, 232], "flatten": [25, 26, 91, 92, 93, 94, 146, 180, 182, 198, 214, 226, 227, 232, 245], "dimension": [27, 112, 113, 116, 117, 118, 119, 120, 121, 125, 126, 127, 258, 259, 260, 261, 262, 266, 273, 274, 275, 299, 307, 388, 390], "val": [27, 131], "tupl": [27, 61, 64, 74, 84, 86, 87, 102, 107, 109, 143, 146, 147, 181, 184, 199, 218, 237, 239, 245, 246, 247, 258, 259, 262, 274, 275, 284, 286, 305, 311, 363, 365, 366, 367, 368, 385, 387], "ndarrai": [27, 388, 389, 391], "properti": [28, 35, 43, 47, 57, 58, 64, 66, 290, 293, 371, 387], "argument": [28, 61, 74, 83, 107, 132, 237, 246, 256, 311, 383, 387, 392, 393, 394], "union": [29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 44, 45, 48, 49, 50, 51, 52, 53, 54, 55, 56, 59, 60, 61, 62, 63, 65, 67, 68, 69, 70, 71, 72, 74, 75, 77, 78, 79, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 140, 141, 194, 195, 222], "appli": [35, 114, 115, 246, 256, 258, 259, 260, 261, 262, 264, 265, 267, 269, 270, 272, 273, 274, 275, 276, 278, 291, 298, 299, 300, 301, 302, 304, 306, 308, 309, 311, 320, 321, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 356, 369, 372, 378, 381, 386], "regular": [35, 264, 343, 366, 386, 388], "idx": [35, 388], "correctli": 35, "syntax": [35, 388], "subtract": 35, "inclus": [37, 38, 39, 40, 91, 92, 93, 94, 128], "diagon": [41, 98, 111, 234, 235, 236], "axis1": [42, 72, 99, 225], "axis2": [42, 72, 99, 225], "start_axi": [45, 128], "end_axi": [45, 128], "datatyp": 47, "byte": [47, 57, 164, 165, 166, 168, 169, 250], "decim": [62, 200], "indices_or_sect": [67, 215], "nest": [73, 83, 256, 359, 385, 387], "ddof": [75, 220, 238], "ari": [77, 78, 79], "a_min": 82, "a_max": 82, "edg": [82, 181, 311, 386], "At": 82, "anoth": [82, 142, 159, 223, 241, 250, 256, 277, 386, 387, 388, 393], "fun": [83, 132, 143, 237, 239, 240, 386, 388, 389, 393], "callabl": [83, 132, 143, 237, 239, 240, 244, 245, 246, 277, 278, 281, 289, 301, 305, 310, 312, 313, 314, 315, 316, 317, 318, 319, 362, 363, 364, 365, 366, 367, 368, 373, 374, 375, 376, 377, 378, 379], "shapeless": 83, "dict": [83, 107, 149, 203, 204, 205, 287, 292, 295, 296, 359, 369, 370, 372, 385, 387, 392], "arbitrarili": [83, 256, 385, 387, 391], "leaf": [83, 245, 246, 281], "node": [83, 107, 240], "recompil": [83, 386], "chang": [83, 208, 295, 299, 311, 334, 341, 386, 390], "Not": [83, 386], "attempt": 83, "pad": [85, 86, 87, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 258, 259, 261, 262, 274, 275], "dilat": [85, 86, 87, 261, 262], "group": [85, 86, 87, 97, 115, 184, 185, 269, 299], "1d": [85, 87, 88, 203, 227], "convolut": [85, 86, 87, 88, 261, 262, 264, 265], "channel": [85, 86, 87, 260, 261, 262, 264, 265], "c_in": [85, 86, 87], "c_out": [85, 86, 87], "convolv": [85, 86, 87], "2d": [86, 87, 99, 184, 260, 264], "spatial": [86, 87, 258, 269, 274, 311], "symmetr": 86, "kernel_dil": 87, "input_dil": 87, "flip": [87, 88], "correl": [87, 264], "discret": [88, 116, 117, 118, 119, 120, 121, 125, 126, 127, 266], "swap": [88, 169, 225, 296, 299], "conv": 88, "filter": [88, 261, 262, 277, 281], "signal": [88, 311], "cumul": [91, 92, 93, 94], "th": [91, 92, 93, 94, 98, 111, 377], "bias": [97, 184, 185, 268, 271, 282, 294, 297], "group_siz": [97, 184, 185, 299], "64": [97, 184, 185, 250, 299], "configur": 97, "formal": [97, 184], "notat": [97, 245, 286], "quantiz": [97, 149, 185, 299], "w_i": [97, 184], "hat": [97, 184], "occupi": [97, 184, 185], "subarrai": [99, 215], "remov": [99, 159, 187, 218, 331], "insert": [99, 109, 393], "neg": [99, 128, 140, 274, 275, 297, 332, 340, 342, 388], "taken": [99, 226], "global": [100, 103, 193, 383, 386], "disabl": [100, 168, 386], "mlx_disable_compil": [100, 103, 386], "divis": [101, 130, 184], "quotient": [101, 102, 130], "remaind": 102, "fuction": 102, "faster": [1, 102, 321, 386, 387], "mathrm": [105, 209, 270], "frac": [105, 184, 209, 258, 259, 260, 263, 264, 265, 269, 270, 272, 273, 274, 275, 300, 313, 314, 315, 316, 330, 332, 334, 337, 348, 350, 362, 364, 365, 366, 367, 373], "pi": [105, 307, 387], "int_0": 105, "dt": 105, "erf": [106, 386], "exponenti": [108, 110, 304, 320, 347, 376], "ident": [111, 221, 256, 291], "zero": [111, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 163, 234, 235, 236, 243, 256, 258, 259, 263, 264, 265, 284, 312, 313, 314, 315, 316, 317, 318, 319, 356, 363, 388], "whose": [111, 244], "translat": [112, 272], "stabil": [112, 113, 260, 269, 270, 272, 300, 330, 332, 362, 363, 364, 365, 366, 367, 373], "traditino": 114, "rotat": [114, 303], "larger": [114, 303, 368], "unchang": [114, 221, 303], "consecut": [114, 184, 303], "angular": [114, 303], "frequenc": [114, 303, 307], "q": [115, 147], "head": [115, 297, 310], "attent": [115, 282, 297, 307, 310], "regardless": 115, "pre": 115, "tile": 115, "typic": [115, 266, 361, 386, 389], "One": [116, 119, 125, 201, 386, 387], "fourier": [116, 117, 118, 119, 120, 121, 125, 126, 127], "truncat": [116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 195], "dft": [116, 117, 118, 119, 120, 121, 125, 126, 127], "rfft": 122, "real": [122, 123, 124, 125, 126, 127], "rfft2": 123, "rfftn": 124, "silent": [125, 126, 127], "outsid": 128, "clamp": 128, "floor": 130, "argnam": [132, 237], "neither": [132, 237], "keyword": [132, 205, 206, 237, 246, 256, 383, 392, 394], "strict": [133, 144, 282, 284, 294], "ordinari": 136, "inifn": 138, "infin": [138, 140, 141, 274, 275, 367], "dtypecategori": [142, 250], "subtyp": [142, 250], "subdtyp": 142, "float64": 142, "too": [142, 386, 389], "ord": 146, "tabl": [146, 250, 266], "frobeniu": 146, "matric": [146, 147], "strictli": 146, "mathemat": 146, "variou": 146, "purpos": 146, "calcul": [146, 332, 338, 363], "fro": 146, "inf": [146, 297], "largest": [146, 232], "sing": 146, "smallest": 146, "singular": 146, "nuclear": 146, "_f": 146, "sum_": [146, 258, 259, 337], "a_": 146, "valueerror": [146, 284, 387], "refer": [146, 270, 276, 290, 313, 314, 315, 316, 323, 343, 388], "golub": 146, "van": 146, "loan": 146, "baltimor": 146, "md": 146, "john": 146, "hopkin": 146, "univers": 146, "1985": 146, "pg": 146, "la": 146, "9": [146, 331, 362, 365, 366, 367, 368, 370, 376, 379, 390], "74597": 146, "84804": 146, "41421": 146, "23607": [146, 147], "74166": 146, "24264": 146, "11": 146, "225": 146, "894427": 147, "447214": 147, "57771": 147, "50": 148, "evenli": 148, "return_metadata": 149, "binari": [149, 202, 203, 204, 205, 206, 309, 329, 353, 386], "npy": [149, 202, 392], "safetensor": [149, 204, 284, 288, 389, 392], "gguf": [149, 203, 392], "matadata": 149, "unsupport": 149, "tensor": [149, 230, 258, 259, 274, 275, 342, 390], "natur": [150, 152, 389], "logarithm": [150, 151, 152, 153], "log": [152, 154, 158, 327, 328, 332, 335, 337, 340, 351], "plu": 152, "exp": [110, 154, 158, 188, 213, 320, 335, 347, 348, 351, 386, 393], "stabl": [154, 158, 213, 337], "prepend": [2, 159], "report": [164, 169], "peak": 166, "begin": [166, 184, 259, 268, 271, 275, 309, 334, 341, 347, 352, 353], "program": 166, "free": 168, "reclaim": 168, "set_memory_limit": 168, "previou": [168, 169], "relax": 169, "task": [169, 337], "exceed": 169, "negat": 176, "beforehand": 180, "pad_with": 181, "constant_valu": 181, "pad_width": 181, "before_1": 181, "after_1": 181, "before_2": 181, "after_2": 181, "before_n": 181, "after_n": 181, "before_i": 181, "after_i": 181, "extend": [1, 181], "side": [181, 258, 259, 274, 275, 386], "smaller": [182, 368, 386], "everi": [184, 246, 379, 387], "particular": [184, 269], "w_1": 184, "w_g": 184, "align": [184, 259, 268, 271, 275], "max_i": 184, "min_i": 184, "textrm": [184, 267, 321, 324], "round": 184, "pack": [184, 185], "unsign": [184, 185, 250], "lower": [184, 192, 195, 196, 234, 319], "upper": [184, 192, 195, 196, 319], "1st": 184, "signific": 184, "2nd": 184, "dequant": 184, "w_q": 184, "whether": [170, 185, 268, 271, 281, 297, 301, 329, 332, 338], "prng": [186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 383], "num_sampl": 187, "unnorm": [187, 329, 331], "draw": 187, "cdf": [188, 267, 321], "accord": [188, 241, 297, 313, 314, 315, 316], "seed": 189, "loc": 191, "deviat": [191, 220, 313, 315, 318], "low": [192, 196, 319, 356], "high": [192, 196, 256, 266, 319, 356], "bound": [192, 195, 196, 267, 319, 386, 388, 393], "roadcast": 192, "domain": 195, "uniformli": 196, "repetit": 198, "preserv": [199, 387], "reciproc": 201, "arr": [202, 388], "obj": 203, "uncompress": 205, "my_path": 205, "tree_flatten": [205, 246, 247, 256], "transformerencod": 205, "128": [205, 256], "flat_param": 205, "compress": 206, "possibl": [215, 266, 386, 388, 393], "being": [221, 256], "prevent": [221, 342, 390], "flow": [221, 389], "streamcontext": 222, "context": 222, "manag": [222, 383, 393], "prior": [226, 227], "exclud": 227, "dot": [230, 245, 286, 297], "rep": 231, "repeat": 231, "necessarili": 232, "elsewher": [234, 388], "col": 234, "triangl": 234, "mse": 237, "param": [237, 256, 356, 387], "lvalu": 237, "dlvalu": 237, "dparam": 237, "lasso": 237, "l1": [237, 334, 336, 337, 341], "varianc": [220, 238, 260, 269, 332], "divisor": [220, 238], "cotang": [1, 239], "in_ax": [240, 387], "out_ax": [240, 387], "prefix": [240, 245], "fn": [244, 246, 391], "wrt": 244, "is_leaf": [245, 246], "arbitrari": [245, 359], "depth": [245, 265, 387], "hello": [245, 247], "charact": 245, "flat": [245, 247], "superset": [246, 369], "extra": 246, "closer": 246, "constitut": 246, "dict_kei": [246, 370], "lambda": [246, 256, 277, 282, 289, 308, 347, 352, 362, 363, 364, 365, 366, 367, 368, 373, 374, 386, 387], "recreat": 247, "world": 247, "42": 247, "16": [250, 258, 270, 274, 277, 359], "int16": 250, "brain": 250, "e8": 250, "m7": 250, "ieee": 250, "e5": 250, "m10": 250, "hierarchi": 250, "done": [256, 263, 300, 386, 389, 390], "manual": 256, "explicitli": [256, 383], "solv": 256, "intuit": 256, "freez": [256, 294, 359], "finetun": 256, "in_dim": [256, 359], "out_dim": [256, 359], "enumer": 256, "caus": [256, 386, 389], "local": [256, 264], "scope": 256, "l2_loss": 256, "y_hat": 256, "trainable_paramet": [256, 281, 370], "loss_and_grad": 256, "workhors": 256, "Its": 256, "recurs": [256, 281, 282, 287, 292, 294, 359], "frozen": [256, 282, 292, 294, 299, 359], "individu": [256, 264, 265], "subset": [256, 281], "action": 256, "displai": 256, "tree_map": 256, "count": [256, 377], "num_param": 256, "preclud": 256, "pure": [256, 361], "pattern": [256, 389], "achiev": 256, "other_input": 256, "necessari": 256, "wrap": 256, "apply_to_modul": [256, 282], "children": 256, "filter_and_map": 256, "leaf_modul": 256, "load_weight": [256, 389], "named_modul": 256, "save_weight": 256, "set_dtyp": 256, "unfreez": [256, 282], "update_modul": 256, "alibi": 256, "avgpool1d": 256, "avgpool2d": 256, "batchnorm": 256, "conv1d": 256, "conv2d": 256, "dropout": [256, 264, 265, 291, 310, 386], "dropout2d": 256, "dropout3d": 256, "gelu": [256, 322, 323, 386], "groupnorm": 256, "gru": 256, "instancenorm": 256, "layernorm": 256, "lstm": 256, "maxpool1d": 256, "maxpool2d": [256, 259], "mish": 256, "prelu": 256, "quantizedlinear": 256, "relu": [256, 298, 310, 344, 356], "rnn": [256, 268], "selu": 256, "sequenti": [256, 356], "silu": 256, "sinusoidalpositionalencod": 256, "softshrink": 256, "upsampl": 256, "elu": [256, 347], "gelu_approx": [256, 267, 321], "gelu_fast_approx": [256, 267, 321], "hardswish": 256, "leaky_relu": 256, "log_sigmoid": 256, "log_softmax": 256, "relu6": 256, "softplu": [256, 276, 343], "tanh": [256, 268, 271, 276, 301, 343], "binary_cross_entropi": [256, 386], "cosine_similarity_loss": 256, "gaussian_nll_loss": 256, "hinge_loss": 256, "huber_loss": 256, "kl_div_loss": 256, "l1_loss": 256, "log_cosh_loss": 256, "margin_ranking_loss": 256, "mse_loss": 256, "nll_loss": 256, "smooth_l1_loss": 256, "triplet_loss": 256, "init": [256, 298, 356, 361, 375, 376, 378, 379], "uniform": [2, 256, 273, 284, 314, 316, 356, 383, 386, 387, 393], "glorot_norm": 256, "glorot_uniform": 256, "he_norm": 256, "he_uniform": 256, "kernel_s": [258, 259, 261, 262, 274, 275], "averag": [258, 259, 362, 363, 365, 366, 367], "pool": [258, 259, 274, 275, 393], "l_": [258, 274, 334], "n_i": [258, 259, 274, 275], "c_j": [258, 259, 274, 275], "ldot": [258, 259, 274, 275], "lfloor": [258, 259, 274, 275], "_size": [258, 259, 274, 275], "rfloor": [258, 259, 274, 275], "k_h": [259, 275], "k_w": [259, 275], "h_": [259, 268, 271, 275, 301], "w_": [259, 268, 271, 275, 301, 362, 363, 364, 365, 366, 367, 368, 373, 374], "height": [259, 260, 262, 264, 265, 275], "width": [259, 260, 262, 264, 265, 275, 299], "momentum": [260, 368, 370, 374, 386], "affin": [260, 269, 270, 272, 273, 299], "track_running_stat": 260, "var": [260, 269, 270, 272, 332], "epsilon": [260, 269, 270, 272, 300, 330, 332, 362, 364, 365, 366, 367, 373], "gamma": [260, 269, 270, 272, 300, 313, 314, 315, 316], "nc": 260, "nlc": [260, 261], "four": 260, "nhwc": [260, 262], "paper": [260, 307, 362, 363, 364, 365, 367, 368], "deep": [260, 313, 314, 315, 316], "intern": 260, "covari": [190, 260], "shift": 260, "bn": 260, "in_channel": [261, 262], "out_channel": [261, 262], "learnabl": [261, 262, 305], "portion": 263, "independ": [264, 265], "nwhc": 264, "whc": 264, "maintain": [264, 265, 368], "entri": [264, 265], "benefici": [264, 265, 389], "earli": 264, "adjac": 264, "pixel": 264, "thompson": 264, "goroshin": 264, "jain": 264, "lecun": 264, "bregler": 264, "2015": [264, 365, 367], "cvpr": 264, "ndhwc": 265, "dhwc": 265, "medic": 265, "video": 265, "num_embed": 266, "lookup": 266, "usual": [266, 385, 389], "vocabulari": 266, "approx": 267, "unit": [267, 268, 302, 304, 306, 313, 314, 315, 316, 320, 321, 322, 323, 324, 326, 345, 346, 347, 349], "phi": [267, 321], "geluapprox": 267, "sigma": [267, 268, 271, 313, 314, 315, 316, 322, 323, 324, 327, 348, 349], "60033": [267, 322], "0433603": [267, 322], "gelufast": 267, "773": 267, "regard": 267, "input_s": [268, 271, 301], "hidden_s": [268, 271, 301], "gate": [268, 324], "recurr": [268, 271, 301], "nld": [268, 271, 301], "ld": [268, 271, 301], "r_t": 268, "xr": 268, "x_t": [268, 271, 301], "hr": 268, "h_t": [268, 271, 301], "b_": [268, 271], "z_t": 268, "xz": 268, "hz": 268, "n_t": 268, "xn": 268, "odot": [268, 271], "hn": 268, "hidden": [268, 271, 301, 310], "nh": [268, 271, 301], "nlh": [268, 271, 301], "lh": [268, 271, 301], "num_group": 269, "pytorch_compat": 269, "split": [269, 324], "preced": 269, "org": [269, 270, 272, 276, 300, 323, 343], "1803": 269, "08494": 269, "denomin": [270, 330, 362, 364, 365, 366, 367, 373], "inorm": 270, "1607": [270, 272], "08022": 270, "i_t": 271, "xi": 271, "f_t": 271, "xf": 271, "hf": 271, "g_t": [271, 362, 364, 365, 366, 367, 368, 373, 374], "xg": 271, "hg": 271, "o_t": 271, "xo": 271, "ho": 271, "c_": [271, 368], "c_t": [271, 368], "cell": 271, "06450": 272, "mathcal": 273, "d_i": 273, "max_": [274, 275], "1908": [276, 343], "08681": [276, 343], "map_fn": [277, 281], "filter_fn": [277, 281], "valid_parameter_filt": 277, "apply_fn": 278, "descend": 279, "is_leaf_fn": 281, "found": 281, "drop": 281, "idempot": [282, 294], "endswith": 282, "file_or_weight": 284, "miss": [284, 392], "ok": [284, 387], "save_safetensor": [288, 392], "predic": 289, "reflect": [290, 386, 388, 390], "certain": [1, 291, 386], "ie": 294, "noop": 294, "unfrozen": 294, "tracer": 295, "partial": [295, 296, 386, 389], "child": 296, "flexibli": 296, "programmat": 296, "query_input_dim": 297, "key_input_dim": 297, "value_input_dim": 297, "value_dim": 297, "value_output_dim": 297, "aggreg": 297, "linearli": 297, "attend": 297, "num_paramet": 298, "25": [298, 311], "parametr": [298, 344], "classmethod": 299, "from_linear": 299, "quantize_modul": 299, "accumul": 300, "1910": 300, "07467": 300, "nonlinear": [301, 386], "elman": 301, "ih": 301, "hh": 301, "func": 301, "rectifi": [302, 315, 316, 326, 345, 346], "10000": 303, "slightli": [303, 393], "plain": 305, "known": [306, 349], "swish": [306, 349], "min_freq": 307, "0001": 307, "max_freq": 307, "cos_first": 307, "full_turn": 307, "sinusoid": 307, "sin": [307, 387, 391], "lambd": [308, 352], "threshold": [309, 334, 341, 353], "geq": [309, 353], "num_encoder_lay": 310, "num_decoder_lay": 310, "nb_func": 310, "custom_encod": 310, "custom_decod": 310, "norm_first": 310, "checkpoint": 310, "decod": 310, "interact": 310, "mechan": 310, "chekpoint": 310, "usag": [310, 386], "expens": 310, "scale_factor": 311, "nearest": 311, "align_corn": 311, "audio": 311, "4d": 311, "forth": 311, "neighbor": 311, "interpol": 311, "bilinear": 311, "trilinear": 311, "corner": 311, "bottom": 311, "squeez": [311, 386], "75": 311, "33333": 311, "66667": 311, "init_fn": [312, 313, 314, 315, 316, 317, 318, 319, 356], "glorot": [313, 314], "fan_in": [313, 314, 315, 316], "fan_out": [313, 314, 315, 316], "fan": [313, 314, 315, 316], "_in": [313, 314], "_out": [313, 314], "difficulti": [313, 314], "feedforward": [313, 314], "191107": 313, "61278": 313, "150594": 313, "363207": 313, "gain": [313, 314, 315, 316], "89613": 313, "53947": 313, "48095": 313, "995016": 313, "223404": 314, "890597": 314, "379159": 314, "776856": 314, "90041": 314, "02264": 314, "912766": 314, "12451": 314, "delv": [315, 316], "surpass": [315, 316], "human": [315, 316], "level": [315, 316], "imagenet": [315, 316], "classif": [315, 316], "25211": 315, "458835": 315, "177208": 315, "0137595": 315, "6967": 315, "02765": 315, "15268": 315, "75787": 315, "kaim": 316, "0300242": 316, "0184009": 316, "793615": 316, "666329": 316, "64331": 316, "16506": 316, "08619": 316, "79854": 316, "982273": 318, "534422": 318, "380709": 318, "0645099": 318, "883935": 319, "863726": 319, "617261": 319, "417497": 319, "exact": [322, 323], "0003": 322, "cdot": [322, 323, 330, 333, 349], "015": 323, "702": 323, "hendryck": 323, "1606": 323, "08415": 323, "halv": 324, "negative_slop": 326, "leaki": 326, "sum_i": 328, "x_i": [328, 350], "with_logit": 329, "predict": [329, 332, 333, 334, 335, 336, 337, 339, 340, 341], "105361": 329, "223144": 329, "20397": 329, "916291": 329, "539245": 329, "prob": 329, "510826": 329, "x1": 330, "x2": 330, "x_1": [330, 338], "x_2": [330, 338], "label_smooth": 331, "hot": 331, "smooth": [331, 341, 373], "0485873": 331, "348587": 331, "06": [332, 342, 362], "likelihood": [332, 340], "nll": [332, 340], "hing": 333, "y_": [333, 337], "pred": [333, 337], "delta": [334, 362], "huber": 334, "leq": [334, 347], "l2": [334, 337, 374], "kullback": 335, "leibler": 335, "diverg": 335, "cosh": 337, "logcosh": 337, "sensit": 337, "outlier": 337, "dual": 337, "behavior": [190, 337, 388, 389], "offer": 337, "balanc": 337, "robust": 337, "approach": [337, 387], "inputs1": 338, "inputs2": 338, "margin": [338, 342], "rank": 338, "573409": 338, "765166": 338, "0638": 338, "75596": 338, "225763": 338, "256995": 338, "773433": 338, "formula": 341, "anchor": 342, "triplet": 342, "_p": 342, "degre": 342, "pairwis": 342, "instabl": 342, "monoton": 343, "0507": 347, "67326": 347, "sum_j": 350, "x_j": 350, "subclass": 359, "concept": 359, "mymlp": 359, "in_proj": 359, "subsequ": 361, "apply_gradi": 361, "rmsprop": 361, "adagrad": 361, "adafactor": 361, "adadelta": 361, "adam": [361, 367, 368, 377, 378], "adamw": [361, 368], "adamax": 361, "lion": 361, "cosine_decai": [361, 377], "exponential_decai": 361, "join_schedul": 361, "linear_schedul": [361, 377], "step_decai": 361, "rho": 362, "zeiler": 362, "2012": [362, 373], "adapt": [362, 363, 364], "1212": 362, "5701": 362, "v_": [362, 364, 365, 366, 367, 373, 374], "v_t": [362, 364, 365, 366, 367, 373, 374], "u_t": 362, "u_": 362, "w_t": [362, 364, 365, 366, 367, 368, 373, 374], "001": 363, "clip_threshold": 363, "decay_r": [363, 376, 379], "beta_1": [363, 365, 366, 367, 368], "weight_decai": [363, 366, 368, 374], "scale_paramet": 363, "relative_step": 363, "warmup_init": 363, "sublinear": 363, "cost": [363, 389], "epsilon_1": 363, "epsilon_2": 363, "parameter_scal": 363, "clip": 363, "unscal": 363, "decai": [363, 366, 368, 374, 375, 376, 379], "duchi": 364, "hazan": 364, "singer": 364, "2011": 364, "subgradi": 364, "onlin": 364, "stochast": [364, 365, 367, 374, 389], "jmlr": 364, "999": [365, 366, 367], "omit": [365, 367], "estim": [365, 367], "kingma": [365, 367], "ba": [365, 367], "iclr": [365, 366, 367], "m_": [365, 366, 367, 368], "m_t": [365, 366, 367, 368], "beta_2": [365, 366, 367, 368], "contrast": 366, "loshchilov": 366, "hutter": 366, "decoupl": 366, "99": [368, 373], "tend": 368, "10x": 368, "strength": [368, 374], "wd": 368, "chen": 368, "symbol": 368, "discoveri": 368, "2302": 368, "06675": 368, "eta": 368, "opt": 369, "tieleman": 373, "hinton": 373, "lectur": 373, "coursera": 373, "dampen": 374, "nesterov": 374, "descent": [374, 386, 389], "mu": 374, "tau": 374, "penalti": 374, "decay_step": 375, "beyond": [375, 378], "minim": 375, "lr_schedul": [375, 376, 377, 379], "1000": [375, 386], "0999961": 375, "06561": 376, "boundari": 377, "join": 377, "receiv": [377, 390], "transit": 377, "warmup": [377, 378], "0999938": 377, "101": 378, "step_siz": 379, "081": 379, "basi": 381, "implicit": [383, 386, 387], "fine": [383, 389], "grain": 383, "control": [383, 389], "pseudo": 383, "altern": 383, "splittabl": 383, "threefri": 383, "counter": 383, "cycl": 385, "merg": 386, "fuse": 386, "big": 386, "awar": [386, 389], "36788": 386, "compiled_fun": 386, "code": [386, 389], "slow": 386, "stack": 386, "rerun": [386, 389], "frequent": [386, 389], "destroi": 386, "anonym": 386, "don": [386, 393], "unari": 386, "overhead": [386, 389, 393], "bandwidth": 386, "fusibl": 386, "consider": 386, "versu": 386, "timeit": [386, 387], "tic": 386, "perf_count": 386, "toc": 386, "tpi": 386, "1e3": 386, "4096": [386, 387, 393], "On": [386, 387, 389], "millisecond": [386, 393], "five": 386, "latest": 386, "won": 386, "placehold": 386, "insid": 386, "crash": 386, "disable_compil": 386, "okai": [386, 389], "intend": 386, "deal": 386, "pretti": [386, 389], "inconveni": 386, "functool": 386, "particularli": 386, "backward": [386, 387], "compiled_grad_fn": 386, "71828": 386, "outer": [386, 389], "opportun": 386, "idea": [387, 389], "behind": 387, "dfdx": [387, 388], "d2fdx2": 387, "zero_grad": 387, "detach": 387, "requires_grad": 387, "dloss_dw": 387, "dloss_dx": 387, "lot": 387, "redund": 387, "suppos": [387, 393], "nice": [387, 389], "propag": [387, 388], "stop_gradi": 387, "autom": 387, "contriv": [387, 393], "sake": 387, "clariti": 387, "quit": [387, 390], "power": [387, 390], "difficult": 387, "primit": 387, "issu": [387, 390], "priorit": 387, "naive_add": 387, "vmap_add": 387, "total": 387, "390": 387, "wherea": 387, "025": 387, "ten": [387, 389], "Of": 387, "better": [387, 393], "handi": 387, "slice": 388, "ellipsi": 388, "mix": 388, "take_along_axi": 388, "lack": 388, "extrem": [388, 389], "ineffici": [388, 389], "nonzero": 388, "dynam": 389, "easier": 389, "worri": 389, "fun1": 389, "expensive_fun": 389, "consum": 389, "eager": 389, "thank": 389, "weights_fp16": 389, "trade": 389, "bad": 389, "grow": 389, "computation": 389, "costli": 389, "wide": 389, "thousand": 389, "value_and_grad_fn": 389, "implicitli": 389, "anytim": 389, "memoryview": [389, 390], "perfectli": 389, "first_lay": 389, "second_layer_a": 389, "second_layer_b": 389, "protocol": 390, "pep": 390, "3118": 390, "a_view": 390, "owndata": 390, "extern": 390, "x_view": 390, "modifi": 390, "df": 390, "x\u00b2": 390, "2x": 390, "indirectli": 390, "modif": 390, "seen": 390, "occur": 390, "incorpor": 390, "incorrect": 390, "experiment": 390, "break": 390, "advis": 390, "intermedi": 390, "jnp": 390, "tf": 390, "page": 391, "composit": 391, "archiv": 392, "savez_compress": 392, "save_gguf": 392, "arr_0": 392, "advantag": 393, "parallel": 393, "race": 393, "interest": 393, "albeit": 393, "d1": 393, "d2": 393, "matmul": 393, "dens": [163, 393], "twice": 393, "measur": 393, "default_stream": 394, "default_devic": 394, "my_devic": 394, "explain": 1, "underli": 1, "tutori": 1, "cover": 1, "front": 1, "virtual": 1, "reimplement": 1, "finish": 1, "command_buff": 1, "mtlcommandbuff": 1, "unus": 1, "nb_modul": 1, "_ext": 1, "nb": 1, "nanobind_add_modul": 1, "nb_static": 1, "stable_abi": 1, "lto": 1, "nomins": 1, "nb_domain": 1, "extras_requir": 1, "cmake_arg": 2, "mtl_capture_en": 2, "trace_fil": 2, "mlx_trace": 2, "exit": 2, "minu": 110, "spars": 163, "xy": 163, "multidimension": 163, "coordin": 163, "cartesian": 163, "ij": 163, "successfulli": 170, "cov": 190, "jointli": 190, "semi": 190, "definit": 190, "empti": 190, "cubic": 311, "bicub": 311}, "objects": {"mlx.core": [[8, 0, 1, "", "Device"], [9, 0, 1, "", "Dtype"], [10, 0, 1, "", "DtypeCategory"], [248, 0, 1, "", "Stream"], [11, 2, 1, "", "abs"], [12, 2, 1, "", "add"], [13, 2, 1, "", "all"], [14, 2, 1, "", "allclose"], [15, 2, 1, "", "any"], [16, 2, 1, "", "arange"], [17, 2, 1, "", "arccos"], [18, 2, 1, "", "arccosh"], [19, 2, 1, "", "arcsin"], [20, 2, 1, "", "arcsinh"], [21, 2, 1, "", "arctan"], [22, 2, 1, "", "arctanh"], [23, 2, 1, "", "argmax"], [24, 2, 1, "", "argmin"], [25, 2, 1, "", "argpartition"], [26, 2, 1, "", "argsort"], [27, 0, 1, "", "array"], [76, 2, 1, "", "array_equal"], [77, 2, 1, "", "atleast_1d"], [78, 2, 1, "", "atleast_2d"], [79, 2, 1, "", "atleast_3d"], [80, 2, 1, "", "broadcast_to"], [81, 2, 1, "", "ceil"], [82, 2, 1, "", "clip"], [83, 2, 1, "", "compile"], [84, 2, 1, "", "concatenate"], [85, 2, 1, "", "conv1d"], [86, 2, 1, "", "conv2d"], [87, 2, 1, "", "conv_general"], [88, 2, 1, "", "convolve"], [89, 2, 1, "", "cos"], [90, 2, 1, "", "cosh"], [91, 2, 1, "", "cummax"], [92, 2, 1, "", "cummin"], [93, 2, 1, "", "cumprod"], [94, 2, 1, "", "cumsum"], [95, 2, 1, "", "default_device"], [96, 2, 1, "", "default_stream"], [97, 2, 1, "", "dequantize"], [98, 2, 1, "", "diag"], [99, 2, 1, "", "diagonal"], [100, 2, 1, "", "disable_compile"], [101, 2, 1, "", "divide"], [102, 2, 1, "", "divmod"], [103, 2, 1, "", "enable_compile"], [104, 2, 1, "", "equal"], [105, 2, 1, "", "erf"], [106, 2, 1, "", "erfinv"], [107, 2, 1, "", "eval"], [108, 2, 1, "", "exp"], [109, 2, 1, "", "expand_dims"], [110, 2, 1, "", "expm1"], [111, 2, 1, "", "eye"], [128, 2, 1, "", "flatten"], [129, 2, 1, "", "floor"], [130, 2, 1, "", "floor_divide"], [131, 2, 1, "", "full"], [132, 2, 1, "", "grad"], [133, 2, 1, "", "greater"], [134, 2, 1, "", "greater_equal"], [135, 2, 1, "", "identity"], [136, 2, 1, "", "inner"], [137, 2, 1, "", "isclose"], [138, 2, 1, "", "isinf"], [139, 2, 1, "", "isnan"], [140, 2, 1, "", "isneginf"], [141, 2, 1, "", "isposinf"], [142, 2, 1, "", "issubdtype"], [143, 2, 1, "", "jvp"], [144, 2, 1, "", "less"], [145, 2, 1, "", "less_equal"], [148, 2, 1, "", "linspace"], [149, 2, 1, "", "load"], [150, 2, 1, "", "log"], [151, 2, 1, "", "log10"], [152, 2, 1, "", "log1p"], [153, 2, 1, "", "log2"], [154, 2, 1, "", "logaddexp"], [155, 2, 1, "", "logical_and"], [156, 2, 1, "", "logical_not"], [157, 2, 1, "", "logical_or"], [158, 2, 1, "", "logsumexp"], [159, 2, 1, "", "matmul"], [160, 2, 1, "", "max"], [161, 2, 1, "", "maximum"], [162, 2, 1, "", "mean"], [163, 2, 1, "", "meshgrid"], [172, 2, 1, "", "min"], [173, 2, 1, "", "minimum"], [174, 2, 1, "", "moveaxis"], [175, 2, 1, "", "multiply"], [176, 2, 1, "", "negative"], [177, 2, 1, "", "new_stream"], [178, 2, 1, "", "ones"], [179, 2, 1, "", "ones_like"], [180, 2, 1, "", "outer"], [181, 2, 1, "", "pad"], [182, 2, 1, "", "partition"], [183, 2, 1, "", "prod"], [184, 2, 1, "", "quantize"], [185, 2, 1, "", "quantized_matmul"], [197, 2, 1, "", "reciprocal"], [198, 2, 1, "", "repeat"], [199, 2, 1, "", "reshape"], [200, 2, 1, "", "round"], [201, 2, 1, "", "rsqrt"], [202, 2, 1, "", "save"], [203, 2, 1, "", "save_gguf"], [204, 2, 1, "", "save_safetensors"], [205, 2, 1, "", "savez"], [206, 2, 1, "", "savez_compressed"], [207, 2, 1, "", "set_default_device"], [208, 2, 1, "", "set_default_stream"], [209, 2, 1, "", "sigmoid"], [210, 2, 1, "", "sign"], [211, 2, 1, "", "sin"], [212, 2, 1, "", "sinh"], [213, 2, 1, "", "softmax"], [214, 2, 1, "", "sort"], [215, 2, 1, "", "split"], [216, 2, 1, "", "sqrt"], [217, 2, 1, "", "square"], [218, 2, 1, "", "squeeze"], [219, 2, 1, "", "stack"], [220, 2, 1, "", "std"], [221, 2, 1, "", "stop_gradient"], [222, 2, 1, "", "stream"], [223, 2, 1, "", "subtract"], [224, 2, 1, "", "sum"], [225, 2, 1, "", "swapaxes"], [226, 2, 1, "", "take"], [227, 2, 1, "", "take_along_axis"], [228, 2, 1, "", "tan"], [229, 2, 1, "", "tanh"], [230, 2, 1, "", "tensordot"], [231, 2, 1, "", "tile"], [232, 2, 1, "", "topk"], [233, 2, 1, "", "transpose"], [234, 2, 1, "", "tri"], [235, 2, 1, "", "tril"], [236, 2, 1, "", "triu"], [237, 2, 1, "", "value_and_grad"], [238, 2, 1, "", "var"], [239, 2, 1, "", "vjp"], [240, 2, 1, "", "vmap"], [241, 2, 1, "", "where"], [242, 2, 1, "", "zeros"], [243, 2, 1, "", "zeros_like"]], "mlx.core.Device": [[8, 1, 1, "", "__init__"]], "mlx.core.Dtype": [[9, 1, 1, "", "__init__"]], "mlx.core.DtypeCategory": [[10, 1, 1, "", "__init__"]], "mlx.core.Stream": [[248, 1, 1, "", "__init__"]], "mlx.core.array": [[28, 3, 1, "", "T"], [27, 1, 1, "", "__init__"], [29, 1, 1, "", "abs"], [30, 1, 1, "", "all"], [31, 1, 1, "", "any"], [32, 1, 1, "", "argmax"], [33, 1, 1, "", "argmin"], [34, 1, 1, "", "astype"], [35, 3, 1, "", "at"], [36, 1, 1, "", "cos"], [37, 1, 1, "", "cummax"], [38, 1, 1, "", "cummin"], [39, 1, 1, "", "cumprod"], [40, 1, 1, "", "cumsum"], [41, 1, 1, "", "diag"], [42, 1, 1, "", "diagonal"], [43, 3, 1, "", "dtype"], [44, 1, 1, "", "exp"], [45, 1, 1, "", "flatten"], [46, 1, 1, "", "item"], [47, 3, 1, "", "itemsize"], [48, 1, 1, "", "log"], [49, 1, 1, "", "log10"], [50, 1, 1, "", "log1p"], [51, 1, 1, "", "log2"], [52, 1, 1, "", "logsumexp"], [53, 1, 1, "", "max"], [54, 1, 1, "", "mean"], [55, 1, 1, "", "min"], [56, 1, 1, "", "moveaxis"], [57, 3, 1, "", "nbytes"], [58, 3, 1, "", "ndim"], [59, 1, 1, "", "prod"], [60, 1, 1, "", "reciprocal"], [61, 1, 1, "", "reshape"], [62, 1, 1, "", "round"], [63, 1, 1, "", "rsqrt"], [64, 3, 1, "", "shape"], [65, 1, 1, "", "sin"], [66, 3, 1, "", "size"], [67, 1, 1, "", "split"], [68, 1, 1, "", "sqrt"], [69, 1, 1, "", "square"], [70, 1, 1, "", "squeeze"], [71, 1, 1, "", "sum"], [72, 1, 1, "", "swapaxes"], [73, 1, 1, "", "tolist"], [74, 1, 1, "", "transpose"], [75, 1, 1, "", "var"]], "mlx.core.fast": [[112, 2, 1, "", "layer_norm"], [113, 2, 1, "", "rms_norm"], [114, 2, 1, "", "rope"], [115, 2, 1, "", "scaled_dot_product_attention"]], "mlx.core.fft": [[116, 2, 1, "", "fft"], [117, 2, 1, "", "fft2"], [118, 2, 1, "", "fftn"], [119, 2, 1, "", "ifft"], [120, 2, 1, "", "ifft2"], [121, 2, 1, "", "ifftn"], [122, 2, 1, "", "irfft"], [123, 2, 1, "", "irfft2"], [124, 2, 1, "", "irfftn"], [125, 2, 1, "", "rfft"], [126, 2, 1, "", "rfft2"], [127, 2, 1, "", "rfftn"]], "mlx.core.linalg": [[146, 2, 1, "", "norm"], [147, 2, 1, "", "qr"]], "mlx.core.metal": [[164, 2, 1, "", "get_active_memory"], [165, 2, 1, "", "get_cache_memory"], [166, 2, 1, "", "get_peak_memory"], [167, 2, 1, "", "is_available"], [168, 2, 1, "", "set_cache_limit"], [169, 2, 1, "", "set_memory_limit"], [170, 2, 1, "", "start_capture"], [171, 2, 1, "", "stop_capture"]], "mlx.core.random": [[186, 2, 1, "", "bernoulli"], [187, 2, 1, "", "categorical"], [188, 2, 1, "", "gumbel"], [189, 2, 1, "", "key"], [190, 2, 1, "", "multivariate_normal"], [191, 2, 1, "", "normal"], [192, 2, 1, "", "randint"], [193, 2, 1, "", "seed"], [194, 2, 1, "", "split"], [195, 2, 1, "", "truncated_normal"], [196, 2, 1, "", "uniform"]], "mlx.nn": [[257, 0, 1, "", "ALiBi"], [258, 0, 1, "", "AvgPool1d"], [259, 0, 1, "", "AvgPool2d"], [260, 0, 1, "", "BatchNorm"], [261, 0, 1, "", "Conv1d"], [262, 0, 1, "", "Conv2d"], [263, 0, 1, "", "Dropout"], [264, 0, 1, "", "Dropout2d"], [265, 0, 1, "", "Dropout3d"], [266, 0, 1, "", "Embedding"], [267, 0, 1, "", "GELU"], [268, 0, 1, "", "GRU"], [269, 0, 1, "", "GroupNorm"], [270, 0, 1, "", "InstanceNorm"], [271, 0, 1, "", "LSTM"], [272, 0, 1, "", "LayerNorm"], [273, 0, 1, "", "Linear"], [274, 0, 1, "", "MaxPool1d"], [275, 0, 1, "", "MaxPool2d"], [276, 0, 1, "", "Mish"], [359, 0, 1, "", "Module"], [297, 0, 1, "", "MultiHeadAttention"], [298, 0, 1, "", "PReLU"], [299, 0, 1, "", "QuantizedLinear"], [300, 0, 1, "", "RMSNorm"], [301, 0, 1, "", "RNN"], [302, 0, 1, "", "ReLU"], [303, 0, 1, "", "RoPE"], [304, 0, 1, "", "SELU"], [305, 0, 1, "", "Sequential"], [306, 0, 1, "", "SiLU"], [307, 0, 1, "", "SinusoidalPositionalEncoding"], [308, 0, 1, "", "Softshrink"], [309, 0, 1, "", "Step"], [310, 0, 1, "", "Transformer"], [311, 0, 1, "", "Upsample"], [320, 2, 1, "", "elu"], [321, 2, 1, "", "gelu"], [322, 2, 1, "", "gelu_approx"], [323, 2, 1, "", "gelu_fast_approx"], [324, 2, 1, "", "glu"], [325, 2, 1, "", "hardswish"], [326, 2, 1, "", "leaky_relu"], [327, 2, 1, "", "log_sigmoid"], [328, 2, 1, "", "log_softmax"], [343, 2, 1, "", "mish"], [344, 2, 1, "", "prelu"], [345, 2, 1, "", "relu"], [346, 2, 1, "", "relu6"], [347, 2, 1, "", "selu"], [348, 2, 1, "", "sigmoid"], [349, 2, 1, "", "silu"], [350, 2, 1, "", "softmax"], [351, 2, 1, "", "softplus"], [352, 2, 1, "", "softshrink"], [353, 2, 1, "", "step"], [354, 2, 1, "", "tanh"], [244, 2, 1, "", "value_and_grad"]], "mlx.nn.Module": [[277, 1, 1, "", "apply"], [278, 1, 1, "", "apply_to_modules"], [279, 1, 1, "", "children"], [280, 1, 1, "", "eval"], [281, 1, 1, "", "filter_and_map"], [282, 1, 1, "", "freeze"], [283, 1, 1, "", "leaf_modules"], [284, 1, 1, "", "load_weights"], [285, 1, 1, "", "modules"], [286, 1, 1, "", "named_modules"], [287, 1, 1, "", "parameters"], [288, 1, 1, "", "save_weights"], [289, 1, 1, "", "set_dtype"], [290, 3, 1, "", "state"], [291, 1, 1, "", "train"], [292, 1, 1, "", "trainable_parameters"], [293, 3, 1, "", "training"], [294, 1, 1, "", "unfreeze"], [295, 1, 1, "", "update"], [296, 1, 1, "", "update_modules"]], "mlx.nn.init": [[312, 2, 1, "", "constant"], [313, 2, 1, "", "glorot_normal"], [314, 2, 1, "", "glorot_uniform"], [315, 2, 1, "", "he_normal"], [316, 2, 1, "", "he_uniform"], [317, 2, 1, "", "identity"], [318, 2, 1, "", "normal"], [319, 2, 1, "", "uniform"]], "mlx.nn.losses": [[329, 2, 1, "", "binary_cross_entropy"], [330, 2, 1, "", "cosine_similarity_loss"], [331, 2, 1, "", "cross_entropy"], [332, 2, 1, "", "gaussian_nll_loss"], [333, 2, 1, "", "hinge_loss"], [334, 2, 1, "", "huber_loss"], [335, 2, 1, "", "kl_div_loss"], [336, 2, 1, "", "l1_loss"], [337, 2, 1, "", "log_cosh_loss"], [338, 2, 1, "", "margin_ranking_loss"], [339, 2, 1, "", "mse_loss"], [340, 2, 1, "", "nll_loss"], [341, 2, 1, "", "smooth_l1_loss"], [342, 2, 1, "", "triplet_loss"]], "mlx.optimizers": [[362, 0, 1, "", "AdaDelta"], [363, 0, 1, "", "Adafactor"], [364, 0, 1, "", "Adagrad"], [365, 0, 1, "", "Adam"], [366, 0, 1, "", "AdamW"], [367, 0, 1, "", "Adamax"], [368, 0, 1, "", "Lion"], [381, 0, 1, "", "Optimizer"], [373, 0, 1, "", "RMSprop"], [374, 0, 1, "", "SGD"], [375, 2, 1, "", "cosine_decay"], [376, 2, 1, "", "exponential_decay"], [377, 2, 1, "", "join_schedules"], [378, 2, 1, "", "linear_schedule"], [379, 2, 1, "", "step_decay"]], "mlx.optimizers.Optimizer": [[369, 1, 1, "", "apply_gradients"], [370, 1, 1, "", "init"], [371, 3, 1, "", "state"], [372, 1, 1, "", "update"]], "mlx.utils": [[245, 2, 1, "", "tree_flatten"], [246, 2, 1, "", "tree_map"], [247, 2, 1, "", "tree_unflatten"]]}, "objtypes": {"0": "py:class", "1": "py:method", "2": "py:function", "3": "py:property"}, "objnames": {"0": ["py", "class", "Python class"], "1": ["py", "method", "Python method"], "2": ["py", "function", "Python function"], "3": ["py", "property", "Python property"]}, "titleterms": {"oper": [0, 1, 360], "develop": 1, "document": 1, "introduc": 1, "exampl": [1, 6, 386, 393], "primit": 1, "us": [1, 389, 394], "implement": [1, 4], "cpu": 1, "backend": [], "gpu": 1, "transform": [1, 310, 384, 386, 387, 389, 391], "build": [1, 7], "bind": 1, "python": [1, 6, 7], "cmake": 1, "setuptool": 1, "usag": [1, 6], "result": 1, "script": [1, 4], "download": [1, 4], "code": [1, 4], "metal": [2, 7, 164, 165, 166, 167, 168, 169, 170, 171, 255], "debugg": 2, "xcode": 2, "workflow": 2, "linear": [3, 254, 273], "regress": 3, "llm": 4, "infer": 4, "model": 4, "attent": 4, "layer": [4, 5, 357], "encod": 4, "full": [4, 131], "gener": 4, "put": 4, "all": [4, 13, 30], "togeth": 4, "convert": 4, "weight": 4, "load": [4, 149, 392], "benchmark": 4, "multi": 5, "perceptron": 5, "mlx": [6, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379], "instal": [6, 7], "api": [6, 7], "refer": 6, "c": [6, 7], "further": 6, "read": 6, "troubleshoot": 7, "from": [7, 388], "sourc": 7, "requir": 7, "option": 7, "found": 7, "x86": 7, "shell": 7, "core": [8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 248], "devic": [8, 251], "dtype": [9, 43], "dtypecategori": 10, "ab": [11, 29], "add": 12, "allclos": 14, "ani": [15, 31], "arang": 16, "arcco": 17, "arccosh": 18, "arcsin": 19, "arcsinh": 20, "arctan": 21, "arctanh": 22, "argmax": [23, 32], "argmin": [24, 33], "argpartit": 25, "argsort": 26, "arrai": [27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 249, 388, 392], "t": 28, "astyp": 34, "co": [36, 89], "cummax": [37, 91], "cummin": [38, 92], "cumprod": [39, 93], "cumsum": [40, 94], "diag": [41, 98], "diagon": [42, 99], "exp": [44, 108], "flatten": [45, 128], "item": 46, "items": 47, "log": [48, 150], "log10": [49, 151], "log1p": [50, 152], "log2": [51, 153], "logsumexp": [52, 158], "max": [53, 160], "mean": [54, 162], "min": [55, 172], "moveaxi": [56, 174], "nbyte": 57, "ndim": 58, "prod": [59, 183], "reciproc": [60, 197], "reshap": [61, 199], "round": [62, 200], "rsqrt": [63, 201], "shape": 64, "sin": [65, 211], "size": 66, "split": [67, 194, 215], "sqrt": [68, 216], "squar": [69, 217], "squeez": [70, 218], "sum": [71, 224], "swapax": [72, 225], "tolist": 73, "transpos": [74, 233], "var": [75, 238], "array_equ": 76, "atleast_1d": 77, "atleast_2d": 78, "atleast_3d": 79, "broadcast_to": 80, "ceil": 81, "clip": 82, "compil": [83, 386], "concaten": 84, "conv1d": [85, 261], "conv2d": [86, 262], "conv_gener": 87, "convolv": 88, "cosh": 90, "default_devic": 95, "default_stream": 96, "dequant": 97, "disable_compil": 100, "divid": 101, "divmod": 102, "enable_compil": 103, "equal": 104, "erf": 105, "erfinv": 106, "eval": [107, 280], "expand_dim": 109, "ey": 111, "fast": [112, 113, 114, 115, 252], "layer_norm": 112, "rms_norm": 113, "rope": [114, 303], "scaled_dot_product_attent": 115, "fft": [116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 253], "fft2": 117, "fftn": 118, "ifft": 119, "ifft2": 120, "ifftn": 121, "irfft": 122, "irfft2": 123, "irfftn": 124, "rfft": 125, "rfft2": 126, "rfftn": 127, "floor": 129, "floor_divid": 130, "grad": [132, 256], "greater": 133, "greater_equ": 134, "ident": [135, 317], "inner": 136, "isclos": 137, "isinf": 138, "isnan": 139, "isneginf": 140, "isposinf": 141, "issubdtyp": 142, "jvp": 143, "less": 144, "less_equ": 145, "linalg": [146, 147], "norm": 146, "qr": 147, "linspac": 148, "logaddexp": 154, "logical_and": 155, "logical_not": 156, "logical_or": 157, "matmul": 159, "maximum": 161, "get_active_memori": 164, "get_cache_memori": 165, "get_peak_memori": 166, "is_avail": 167, "set_cache_limit": 168, "set_memory_limit": 169, "minimum": 173, "multipli": 175, "neg": 176, "new_stream": 177, "ones": 178, "ones_lik": 179, "outer": 180, "pad": 181, "partit": 182, "quantiz": 184, "quantized_matmul": 185, "random": [186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 383], "bernoulli": 186, "categor": 187, "gumbel": 188, "kei": 189, "normal": [191, 318], "randint": 192, "seed": 193, "truncated_norm": 195, "uniform": [196, 319], "repeat": 198, "save": [202, 392], "save_gguf": 203, "save_safetensor": 204, "savez": 205, "savez_compress": 206, "set_default_devic": 207, "set_default_stream": 208, "sigmoid": [209, 348], "sign": 210, "sinh": 212, "softmax": [213, 350], "sort": 214, "stack": 219, "stop_gradi": 221, "stream": [222, 248, 251, 394], "subtract": 223, "take": 226, "take_along_axi": 227, "tan": 228, "tanh": [229, 354], "tensordot": 230, "tile": 231, "topk": 232, "tri": 234, "tril": 235, "triu": 236, "value_and_grad": [237, 244], "vjp": 239, "vmap": 240, "where": 241, "zero": 242, "zeros_lik": 243, "nn": [244, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354], "util": [245, 246, 247, 385], "tree_flatten": 245, "tree_map": 246, "tree_unflatten": 247, "data": 250, "type": 250, "support": 250, "algebra": 254, "neural": 256, "network": 256, "quick": [256, 391], "start": [256, 391], "The": 256, "modul": [256, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 359], "class": 256, "paramet": [256, 287], "updat": [256, 295, 372, 388], "inspect": 256, "valu": 256, "alibi": 257, "avgpool1d": 258, "avgpool2d": 259, "batchnorm": 260, "dropout": 263, "dropout2d": 264, "dropout3d": 265, "embed": 266, "gelu": [267, 321], "gru": 268, "groupnorm": 269, "instancenorm": 270, "lstm": 271, "layernorm": 272, "maxpool1d": 274, "maxpool2d": 275, "mish": [276, 343], "appli": 277, "apply_to_modul": 278, "children": 279, "filter_and_map": 281, "freez": 282, "leaf_modul": 283, "load_weight": 284, "named_modul": 286, "save_weight": 288, "set_dtyp": 289, "state": [290, 371], "train": [291, 293, 386], "trainable_paramet": 292, "unfreez": 294, "update_modul": 296, "multiheadattent": 297, "prelu": [298, 344], "quantizedlinear": 299, "rmsnorm": 300, "rnn": 301, "relu": [302, 345], "selu": [304, 347], "sequenti": 305, "silu": [306, 349], "sinusoidalpositionalencod": 307, "softshrink": [308, 352], "step": [309, 353], "upsampl": 311, "init": [312, 313, 314, 315, 316, 317, 318, 319, 370], "constant": 312, "glorot_norm": 313, "glorot_uniform": 314, "he_norm": 315, "he_uniform": 316, "elu": 320, "gelu_approx": 322, "gelu_fast_approx": 323, "glu": 324, "hardswish": 325, "leaky_relu": 326, "log_sigmoid": 327, "log_softmax": 328, "loss": [329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 358], "binary_cross_entropi": 329, "cosine_similarity_loss": 330, "cross_entropi": 331, "gaussian_nll_loss": 332, "hinge_loss": 333, "huber_loss": 334, "kl_div_loss": 335, "l1_loss": 336, "log_cosh_loss": 337, "margin_ranking_loss": 338, "mse_loss": 339, "nll_loss": 340, "smooth_l1_loss": 341, "triplet_loss": 342, "relu6": 346, "softplu": 351, "function": [355, 358, 386, 387, 391], "initi": 356, "optim": [361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381], "adadelta": 362, "adafactor": 363, "adagrad": 364, "adam": 365, "adamw": 366, "adamax": 367, "lion": 368, "apply_gradi": 369, "rmsprop": 373, "sgd": 374, "cosine_decai": 375, "exponential_decai": 376, "join_schedul": 377, "linear_schedul": 378, "step_decai": 379, "common": 380, "schedul": 382, "tree": 385, "basic": [386, 391], "speedup": 386, "debug": 386, "pure": 386, "graph": [386, 389, 391], "automat": 387, "differenti": 387, "vector": 387, "index": 388, "differ": 388, "numpi": [388, 390], "In": 388, "place": 388, "lazi": 389, "evalu": 389, "why": 389, "comput": 389, "onli": 389, "what": 389, "you": 389, "when": 389, "convers": 390, "other": 390, "framework": 390, "pytorch": 390, "jax": 390, "tensorflow": 390, "guid": 391, "serial": 392, "format": 392, "unifi": 393, "memori": 393, "A": 393, "simpl": 393, "specifi": 394, "back": 1, "end": 1, "expm1": 110, "meshgrid": 163, "start_captur": 170, "stop_captur": 171, "multivariate_norm": 190, "std": 220}, "envversion": {"sphinx.domains.c": 3, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 9, "sphinx.domains.index": 1, "sphinx.domains.javascript": 3, "sphinx.domains.math": 2, "sphinx.domains.python": 4, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.intersphinx": 1, "sphinx": 60}, "alltitles": {"Operations": [[0, "operations"], [1, "operations"], [360, "operations"]], "Developer Documentation": [[1, "developer-documentation"]], "Introducing the Example": [[1, "introducing-the-example"]], "Operations and Primitives": [[1, "operations-and-primitives"]], "Primitives": [[1, "primitives"]], "Using the Primitive": [[1, "using-the-primitive"]], "Implementing the Primitive": [[1, "implementing-the-primitive"]], "Implementing the CPU Back-end": [[1, "implementing-the-cpu-back-end"]], "Implementing the GPU Back-end": [[1, "implementing-the-gpu-back-end"]], "Primitive Transforms": [[1, "primitive-transforms"]], "Building and Binding": [[1, "building-and-binding"]], "Binding to Python": [[1, "binding-to-python"]], "Building with CMake": [[1, "building-with-cmake"]], "Building with setuptools": [[1, "building-with-setuptools"]], "Usage": [[1, "usage"], [6, null]], "Results": [[1, "results"]], "Scripts": [[1, "scripts"], [4, "scripts"]], "Download the code": [[1, null], [4, null]], "Metal Debugger": [[2, "metal-debugger"]], "Xcode Workflow": [[2, "xcode-workflow"]], "Linear Regression": [[3, "linear-regression"]], "LLM inference": [[4, "llm-inference"]], "Implementing the model": [[4, "implementing-the-model"]], "Attention layer": [[4, "attention-layer"]], "Encoder layer": [[4, "encoder-layer"]], "Full model": [[4, "full-model"]], "Generation": [[4, "generation"]], "Putting it all together": [[4, "putting-it-all-together"]], "Converting the weights": [[4, "converting-the-weights"]], "Weight loading and benchmarking": [[4, "weight-loading-and-benchmarking"]], "Multi-Layer Perceptron": [[5, "multi-layer-perceptron"]], "MLX": [[6, "mlx"]], "Install": [[6, null]], "Examples": [[6, null]], "Python API Reference": [[6, null]], "C++ API Reference": [[6, null]], "Further Reading": [[6, null]], "Build and Install": [[7, "build-and-install"]], "Python Installation": [[7, "python-installation"]], "Troubleshooting": [[7, "troubleshooting"], [7, "id2"]], "Build from source": [[7, "build-from-source"]], "Build Requirements": [[7, "build-requirements"]], "Python API": [[7, "python-api"]], "C++ API": [[7, "c-api"]], "Build Options": [[7, "id3"]], "Metal not found": [[7, "metal-not-found"]], "x86 Shell": [[7, "x86-shell"]], "mlx.core.Device": [[8, "mlx-core-device"]], "mlx.core.Dtype": [[9, "mlx-core-dtype"]], "mlx.core.DtypeCategory": [[10, "mlx-core-dtypecategory"]], "mlx.core.abs": [[11, "mlx-core-abs"]], "mlx.core.add": [[12, "mlx-core-add"]], "mlx.core.all": [[13, "mlx-core-all"]], "mlx.core.allclose": [[14, "mlx-core-allclose"]], "mlx.core.any": [[15, "mlx-core-any"]], "mlx.core.arange": [[16, "mlx-core-arange"]], "mlx.core.arccos": [[17, "mlx-core-arccos"]], "mlx.core.arccosh": [[18, "mlx-core-arccosh"]], "mlx.core.arcsin": [[19, "mlx-core-arcsin"]], "mlx.core.arcsinh": [[20, "mlx-core-arcsinh"]], "mlx.core.arctan": [[21, "mlx-core-arctan"]], "mlx.core.arctanh": [[22, "mlx-core-arctanh"]], "mlx.core.argmax": [[23, "mlx-core-argmax"]], "mlx.core.argmin": [[24, "mlx-core-argmin"]], "mlx.core.argpartition": [[25, "mlx-core-argpartition"]], "mlx.core.argsort": [[26, "mlx-core-argsort"]], "mlx.core.array": [[27, "mlx-core-array"]], "mlx.core.array.T": [[28, "mlx-core-array-t"]], "mlx.core.array.abs": [[29, "mlx-core-array-abs"]], "mlx.core.array.all": [[30, "mlx-core-array-all"]], "mlx.core.array.any": [[31, "mlx-core-array-any"]], "mlx.core.array.argmax": [[32, "mlx-core-array-argmax"]], "mlx.core.array.argmin": [[33, "mlx-core-array-argmin"]], "mlx.core.array.astype": [[34, "mlx-core-array-astype"]], "mlx.core.array.at": [[35, "mlx-core-array-at"]], "mlx.core.array.cos": [[36, "mlx-core-array-cos"]], "mlx.core.array.cummax": [[37, "mlx-core-array-cummax"]], "mlx.core.array.cummin": [[38, "mlx-core-array-cummin"]], "mlx.core.array.cumprod": [[39, "mlx-core-array-cumprod"]], "mlx.core.array.cumsum": [[40, "mlx-core-array-cumsum"]], "mlx.core.array.diag": [[41, "mlx-core-array-diag"]], "mlx.core.array.diagonal": [[42, "mlx-core-array-diagonal"]], "mlx.core.array.dtype": [[43, "mlx-core-array-dtype"]], "mlx.core.array.exp": [[44, "mlx-core-array-exp"]], "mlx.core.array.flatten": [[45, "mlx-core-array-flatten"]], "mlx.core.array.item": [[46, "mlx-core-array-item"]], "mlx.core.array.itemsize": [[47, "mlx-core-array-itemsize"]], "mlx.core.array.log": [[48, "mlx-core-array-log"]], "mlx.core.array.log10": [[49, "mlx-core-array-log10"]], "mlx.core.array.log1p": [[50, "mlx-core-array-log1p"]], "mlx.core.array.log2": [[51, "mlx-core-array-log2"]], "mlx.core.array.logsumexp": [[52, "mlx-core-array-logsumexp"]], "mlx.core.array.max": [[53, "mlx-core-array-max"]], "mlx.core.array.mean": [[54, "mlx-core-array-mean"]], "mlx.core.array.min": [[55, "mlx-core-array-min"]], "mlx.core.array.moveaxis": [[56, "mlx-core-array-moveaxis"]], "mlx.core.array.nbytes": [[57, "mlx-core-array-nbytes"]], "mlx.core.array.ndim": [[58, "mlx-core-array-ndim"]], "mlx.core.array.prod": [[59, "mlx-core-array-prod"]], "mlx.core.array.reciprocal": [[60, "mlx-core-array-reciprocal"]], "mlx.core.array.reshape": [[61, "mlx-core-array-reshape"]], "mlx.core.array.round": [[62, "mlx-core-array-round"]], "mlx.core.array.rsqrt": [[63, "mlx-core-array-rsqrt"]], "mlx.core.array.shape": [[64, "mlx-core-array-shape"]], "mlx.core.array.sin": [[65, "mlx-core-array-sin"]], "mlx.core.array.size": [[66, "mlx-core-array-size"]], "mlx.core.array.split": [[67, "mlx-core-array-split"]], "mlx.core.array.sqrt": [[68, "mlx-core-array-sqrt"]], "mlx.core.array.square": [[69, "mlx-core-array-square"]], "mlx.core.array.squeeze": [[70, "mlx-core-array-squeeze"]], "mlx.core.array.sum": [[71, "mlx-core-array-sum"]], "mlx.core.array.swapaxes": [[72, "mlx-core-array-swapaxes"]], "mlx.core.array.tolist": [[73, "mlx-core-array-tolist"]], "mlx.core.array.transpose": [[74, "mlx-core-array-transpose"]], "mlx.core.array.var": [[75, "mlx-core-array-var"]], "mlx.core.array_equal": [[76, "mlx-core-array-equal"]], "mlx.core.atleast_1d": [[77, "mlx-core-atleast-1d"]], "mlx.core.atleast_2d": [[78, "mlx-core-atleast-2d"]], "mlx.core.atleast_3d": [[79, "mlx-core-atleast-3d"]], "mlx.core.broadcast_to": [[80, "mlx-core-broadcast-to"]], "mlx.core.ceil": [[81, "mlx-core-ceil"]], "mlx.core.clip": [[82, "mlx-core-clip"]], "mlx.core.compile": [[83, "mlx-core-compile"]], "mlx.core.concatenate": [[84, "mlx-core-concatenate"]], "mlx.core.conv1d": [[85, "mlx-core-conv1d"]], "mlx.core.conv2d": [[86, "mlx-core-conv2d"]], "mlx.core.conv_general": [[87, "mlx-core-conv-general"]], "mlx.core.convolve": [[88, "mlx-core-convolve"]], "mlx.core.cos": [[89, "mlx-core-cos"]], "mlx.core.cosh": [[90, "mlx-core-cosh"]], "mlx.core.cummax": [[91, "mlx-core-cummax"]], "mlx.core.cummin": [[92, "mlx-core-cummin"]], "mlx.core.cumprod": [[93, "mlx-core-cumprod"]], "mlx.core.cumsum": [[94, "mlx-core-cumsum"]], "mlx.core.default_device": [[95, "mlx-core-default-device"]], "mlx.core.default_stream": [[96, "mlx-core-default-stream"]], "mlx.core.dequantize": [[97, "mlx-core-dequantize"]], "mlx.core.diag": [[98, "mlx-core-diag"]], "mlx.core.diagonal": [[99, "mlx-core-diagonal"]], "mlx.core.disable_compile": [[100, "mlx-core-disable-compile"]], "mlx.core.divide": [[101, "mlx-core-divide"]], "mlx.core.divmod": [[102, "mlx-core-divmod"]], "mlx.core.enable_compile": [[103, "mlx-core-enable-compile"]], "mlx.core.equal": [[104, "mlx-core-equal"]], "mlx.core.erf": [[105, "mlx-core-erf"]], "mlx.core.erfinv": [[106, "mlx-core-erfinv"]], "mlx.core.eval": [[107, "mlx-core-eval"]], "mlx.core.exp": [[108, "mlx-core-exp"]], "mlx.core.expand_dims": [[109, "mlx-core-expand-dims"]], "mlx.core.expm1": [[110, "mlx-core-expm1"]], "mlx.core.eye": [[111, "mlx-core-eye"]], "mlx.core.fast.layer_norm": [[112, "mlx-core-fast-layer-norm"]], "mlx.core.fast.rms_norm": [[113, "mlx-core-fast-rms-norm"]], "mlx.core.fast.rope": [[114, "mlx-core-fast-rope"]], "mlx.core.fast.scaled_dot_product_attention": [[115, "mlx-core-fast-scaled-dot-product-attention"]], "mlx.core.fft.fft": [[116, "mlx-core-fft-fft"]], "mlx.core.fft.fft2": [[117, "mlx-core-fft-fft2"]], "mlx.core.fft.fftn": [[118, "mlx-core-fft-fftn"]], "mlx.core.fft.ifft": [[119, "mlx-core-fft-ifft"]], "mlx.core.fft.ifft2": [[120, "mlx-core-fft-ifft2"]], "mlx.core.fft.ifftn": [[121, "mlx-core-fft-ifftn"]], "mlx.core.fft.irfft": [[122, "mlx-core-fft-irfft"]], "mlx.core.fft.irfft2": [[123, "mlx-core-fft-irfft2"]], "mlx.core.fft.irfftn": [[124, "mlx-core-fft-irfftn"]], "mlx.core.fft.rfft": [[125, "mlx-core-fft-rfft"]], "mlx.core.fft.rfft2": [[126, "mlx-core-fft-rfft2"]], "mlx.core.fft.rfftn": [[127, "mlx-core-fft-rfftn"]], "mlx.core.flatten": [[128, "mlx-core-flatten"]], "mlx.core.floor": [[129, "mlx-core-floor"]], "mlx.core.floor_divide": [[130, "mlx-core-floor-divide"]], "mlx.core.full": [[131, "mlx-core-full"]], "mlx.core.grad": [[132, "mlx-core-grad"]], "mlx.core.greater": [[133, "mlx-core-greater"]], "mlx.core.greater_equal": [[134, "mlx-core-greater-equal"]], "mlx.core.identity": [[135, "mlx-core-identity"]], "mlx.core.inner": [[136, "mlx-core-inner"]], "mlx.core.isclose": [[137, "mlx-core-isclose"]], "mlx.core.isinf": [[138, "mlx-core-isinf"]], "mlx.core.isnan": [[139, "mlx-core-isnan"]], "mlx.core.isneginf": [[140, "mlx-core-isneginf"]], "mlx.core.isposinf": [[141, "mlx-core-isposinf"]], "mlx.core.issubdtype": [[142, "mlx-core-issubdtype"]], "mlx.core.jvp": [[143, "mlx-core-jvp"]], "mlx.core.less": [[144, "mlx-core-less"]], "mlx.core.less_equal": [[145, "mlx-core-less-equal"]], "mlx.core.linalg.norm": [[146, "mlx-core-linalg-norm"]], "mlx.core.linalg.qr": [[147, "mlx-core-linalg-qr"]], "mlx.core.linspace": [[148, "mlx-core-linspace"]], "mlx.core.load": [[149, "mlx-core-load"]], "mlx.core.log": [[150, "mlx-core-log"]], "mlx.core.log10": [[151, "mlx-core-log10"]], "mlx.core.log1p": [[152, "mlx-core-log1p"]], "mlx.core.log2": [[153, "mlx-core-log2"]], "mlx.core.logaddexp": [[154, "mlx-core-logaddexp"]], "mlx.core.logical_and": [[155, "mlx-core-logical-and"]], "mlx.core.logical_not": [[156, "mlx-core-logical-not"]], "mlx.core.logical_or": [[157, "mlx-core-logical-or"]], "mlx.core.logsumexp": [[158, "mlx-core-logsumexp"]], "mlx.core.matmul": [[159, "mlx-core-matmul"]], "mlx.core.max": [[160, "mlx-core-max"]], "mlx.core.maximum": [[161, "mlx-core-maximum"]], "mlx.core.mean": [[162, "mlx-core-mean"]], "mlx.core.meshgrid": [[163, "mlx-core-meshgrid"]], "mlx.core.metal.get_active_memory": [[164, "mlx-core-metal-get-active-memory"]], "mlx.core.metal.get_cache_memory": [[165, "mlx-core-metal-get-cache-memory"]], "mlx.core.metal.get_peak_memory": [[166, "mlx-core-metal-get-peak-memory"]], "mlx.core.metal.is_available": [[167, "mlx-core-metal-is-available"]], "mlx.core.metal.set_cache_limit": [[168, "mlx-core-metal-set-cache-limit"]], "mlx.core.metal.set_memory_limit": [[169, "mlx-core-metal-set-memory-limit"]], "mlx.core.metal.start_capture": [[170, "mlx-core-metal-start-capture"]], "mlx.core.metal.stop_capture": [[171, "mlx-core-metal-stop-capture"]], "mlx.core.min": [[172, "mlx-core-min"]], "mlx.core.minimum": [[173, "mlx-core-minimum"]], "mlx.core.moveaxis": [[174, "mlx-core-moveaxis"]], "mlx.core.multiply": [[175, "mlx-core-multiply"]], "mlx.core.negative": [[176, "mlx-core-negative"]], "mlx.core.new_stream": [[177, "mlx-core-new-stream"]], "mlx.core.ones": [[178, "mlx-core-ones"]], "mlx.core.ones_like": [[179, "mlx-core-ones-like"]], "mlx.core.outer": [[180, "mlx-core-outer"]], "mlx.core.pad": [[181, "mlx-core-pad"]], "mlx.core.partition": [[182, "mlx-core-partition"]], "mlx.core.prod": [[183, "mlx-core-prod"]], "mlx.core.quantize": [[184, "mlx-core-quantize"]], "mlx.core.quantized_matmul": [[185, "mlx-core-quantized-matmul"]], "mlx.core.random.bernoulli": [[186, "mlx-core-random-bernoulli"]], "mlx.core.random.categorical": [[187, "mlx-core-random-categorical"]], "mlx.core.random.gumbel": [[188, "mlx-core-random-gumbel"]], "mlx.core.random.key": [[189, "mlx-core-random-key"]], "mlx.core.random.multivariate_normal": [[190, "mlx-core-random-multivariate-normal"]], "mlx.core.random.normal": [[191, "mlx-core-random-normal"]], "mlx.core.random.randint": [[192, "mlx-core-random-randint"]], "mlx.core.random.seed": [[193, "mlx-core-random-seed"]], "mlx.core.random.split": [[194, "mlx-core-random-split"]], "mlx.core.random.truncated_normal": [[195, "mlx-core-random-truncated-normal"]], "mlx.core.random.uniform": [[196, "mlx-core-random-uniform"]], "mlx.core.reciprocal": [[197, "mlx-core-reciprocal"]], "mlx.core.repeat": [[198, "mlx-core-repeat"]], "mlx.core.reshape": [[199, "mlx-core-reshape"]], "mlx.core.round": [[200, "mlx-core-round"]], "mlx.core.rsqrt": [[201, "mlx-core-rsqrt"]], "mlx.core.save": [[202, "mlx-core-save"]], "mlx.core.save_gguf": [[203, "mlx-core-save-gguf"]], "mlx.core.save_safetensors": [[204, "mlx-core-save-safetensors"]], "mlx.core.savez": [[205, "mlx-core-savez"]], "mlx.core.savez_compressed": [[206, "mlx-core-savez-compressed"]], "mlx.core.set_default_device": [[207, "mlx-core-set-default-device"]], "mlx.core.set_default_stream": [[208, "mlx-core-set-default-stream"]], "mlx.core.sigmoid": [[209, "mlx-core-sigmoid"]], "mlx.core.sign": [[210, "mlx-core-sign"]], "mlx.core.sin": [[211, "mlx-core-sin"]], "mlx.core.sinh": [[212, "mlx-core-sinh"]], "mlx.core.softmax": [[213, "mlx-core-softmax"]], "mlx.core.sort": [[214, "mlx-core-sort"]], "mlx.core.split": [[215, "mlx-core-split"]], "mlx.core.sqrt": [[216, "mlx-core-sqrt"]], "mlx.core.square": [[217, "mlx-core-square"]], "mlx.core.squeeze": [[218, "mlx-core-squeeze"]], "mlx.core.stack": [[219, "mlx-core-stack"]], "mlx.core.std": [[220, "mlx-core-std"]], "mlx.core.stop_gradient": [[221, "mlx-core-stop-gradient"]], "mlx.core.stream": [[222, "mlx-core-stream"]], "mlx.core.subtract": [[223, "mlx-core-subtract"]], "mlx.core.sum": [[224, "mlx-core-sum"]], "mlx.core.swapaxes": [[225, "mlx-core-swapaxes"]], "mlx.core.take": [[226, "mlx-core-take"]], "mlx.core.take_along_axis": [[227, "mlx-core-take-along-axis"]], "mlx.core.tan": [[228, "mlx-core-tan"]], "mlx.core.tanh": [[229, "mlx-core-tanh"]], "mlx.core.tensordot": [[230, "mlx-core-tensordot"]], "mlx.core.tile": [[231, "mlx-core-tile"]], "mlx.core.topk": [[232, "mlx-core-topk"]], "mlx.core.transpose": [[233, "mlx-core-transpose"]], "mlx.core.tri": [[234, "mlx-core-tri"]], "mlx.core.tril": [[235, "mlx-core-tril"]], "mlx.core.triu": [[236, "mlx-core-triu"]], "mlx.core.value_and_grad": [[237, "mlx-core-value-and-grad"]], "mlx.core.var": [[238, "mlx-core-var"]], "mlx.core.vjp": [[239, "mlx-core-vjp"]], "mlx.core.vmap": [[240, "mlx-core-vmap"]], "mlx.core.where": [[241, "mlx-core-where"]], "mlx.core.zeros": [[242, "mlx-core-zeros"]], "mlx.core.zeros_like": [[243, "mlx-core-zeros-like"]], "mlx.nn.value_and_grad": [[244, "mlx-nn-value-and-grad"]], "mlx.utils.tree_flatten": [[245, "mlx-utils-tree-flatten"]], "mlx.utils.tree_map": [[246, "mlx-utils-tree-map"]], "mlx.utils.tree_unflatten": [[247, "mlx-utils-tree-unflatten"]], "mlx.core.Stream": [[248, "mlx-core-stream"]], "Array": [[249, "array"]], "Data Types": [[250, "data-types"]], "Supported Data Types": [[250, "id2"]], "Devices and Streams": [[251, "devices-and-streams"]], "Fast": [[252, "fast"]], "FFT": [[253, "fft"]], "Linear Algebra": [[254, "linear-algebra"]], "Metal": [[255, "metal"]], "Neural Networks": [[256, "neural-networks"]], "Quick Start with Neural Networks": [[256, "quick-start-with-neural-networks"]], "The Module Class": [[256, "the-module-class"]], "Parameters": [[256, "parameters"]], "Updating the Parameters": [[256, "updating-the-parameters"]], "Inspecting Modules": [[256, "inspecting-modules"]], "Value and Grad": [[256, "value-and-grad"]], "mlx.nn.ALiBi": [[257, "mlx-nn-alibi"]], "mlx.nn.AvgPool1d": [[258, "mlx-nn-avgpool1d"]], "mlx.nn.AvgPool2d": [[259, "mlx-nn-avgpool2d"]], "mlx.nn.BatchNorm": [[260, "mlx-nn-batchnorm"]], "mlx.nn.Conv1d": [[261, "mlx-nn-conv1d"]], "mlx.nn.Conv2d": [[262, "mlx-nn-conv2d"]], "mlx.nn.Dropout": [[263, "mlx-nn-dropout"]], "mlx.nn.Dropout2d": [[264, "mlx-nn-dropout2d"]], "mlx.nn.Dropout3d": [[265, "mlx-nn-dropout3d"]], "mlx.nn.Embedding": [[266, "mlx-nn-embedding"]], "mlx.nn.GELU": [[267, "mlx-nn-gelu"]], "mlx.nn.GRU": [[268, "mlx-nn-gru"]], "mlx.nn.GroupNorm": [[269, "mlx-nn-groupnorm"]], "mlx.nn.InstanceNorm": [[270, "mlx-nn-instancenorm"]], "mlx.nn.LSTM": [[271, "mlx-nn-lstm"]], "mlx.nn.LayerNorm": [[272, "mlx-nn-layernorm"]], "mlx.nn.Linear": [[273, "mlx-nn-linear"]], "mlx.nn.MaxPool1d": [[274, "mlx-nn-maxpool1d"]], "mlx.nn.MaxPool2d": [[275, "mlx-nn-maxpool2d"]], "mlx.nn.Mish": [[276, "mlx-nn-mish"]], "mlx.nn.Module.apply": [[277, "mlx-nn-module-apply"]], "mlx.nn.Module.apply_to_modules": [[278, "mlx-nn-module-apply-to-modules"]], "mlx.nn.Module.children": [[279, "mlx-nn-module-children"]], "mlx.nn.Module.eval": [[280, "mlx-nn-module-eval"]], "mlx.nn.Module.filter_and_map": [[281, "mlx-nn-module-filter-and-map"]], "mlx.nn.Module.freeze": [[282, "mlx-nn-module-freeze"]], "mlx.nn.Module.leaf_modules": [[283, "mlx-nn-module-leaf-modules"]], "mlx.nn.Module.load_weights": [[284, "mlx-nn-module-load-weights"]], "mlx.nn.Module.modules": [[285, "mlx-nn-module-modules"]], "mlx.nn.Module.named_modules": [[286, "mlx-nn-module-named-modules"]], "mlx.nn.Module.parameters": [[287, "mlx-nn-module-parameters"]], "mlx.nn.Module.save_weights": [[288, "mlx-nn-module-save-weights"]], "mlx.nn.Module.set_dtype": [[289, "mlx-nn-module-set-dtype"]], "mlx.nn.Module.state": [[290, "mlx-nn-module-state"]], "mlx.nn.Module.train": [[291, "mlx-nn-module-train"]], "mlx.nn.Module.trainable_parameters": [[292, "mlx-nn-module-trainable-parameters"]], "mlx.nn.Module.training": [[293, "mlx-nn-module-training"]], "mlx.nn.Module.unfreeze": [[294, "mlx-nn-module-unfreeze"]], "mlx.nn.Module.update": [[295, "mlx-nn-module-update"]], "mlx.nn.Module.update_modules": [[296, "mlx-nn-module-update-modules"]], "mlx.nn.MultiHeadAttention": [[297, "mlx-nn-multiheadattention"]], "mlx.nn.PReLU": [[298, "mlx-nn-prelu"]], "mlx.nn.QuantizedLinear": [[299, "mlx-nn-quantizedlinear"]], "mlx.nn.RMSNorm": [[300, "mlx-nn-rmsnorm"]], "mlx.nn.RNN": [[301, "mlx-nn-rnn"]], "mlx.nn.ReLU": [[302, "mlx-nn-relu"]], "mlx.nn.RoPE": [[303, "mlx-nn-rope"]], "mlx.nn.SELU": [[304, "mlx-nn-selu"]], "mlx.nn.Sequential": [[305, "mlx-nn-sequential"]], "mlx.nn.SiLU": [[306, "mlx-nn-silu"]], "mlx.nn.SinusoidalPositionalEncoding": [[307, "mlx-nn-sinusoidalpositionalencoding"]], "mlx.nn.Softshrink": [[308, "mlx-nn-softshrink"]], "mlx.nn.Step": [[309, "mlx-nn-step"]], "mlx.nn.Transformer": [[310, "mlx-nn-transformer"]], "mlx.nn.Upsample": [[311, "mlx-nn-upsample"]], "mlx.nn.init.constant": [[312, "mlx-nn-init-constant"]], "mlx.nn.init.glorot_normal": [[313, "mlx-nn-init-glorot-normal"]], "mlx.nn.init.glorot_uniform": [[314, "mlx-nn-init-glorot-uniform"]], "mlx.nn.init.he_normal": [[315, "mlx-nn-init-he-normal"]], "mlx.nn.init.he_uniform": [[316, "mlx-nn-init-he-uniform"]], "mlx.nn.init.identity": [[317, "mlx-nn-init-identity"]], "mlx.nn.init.normal": [[318, "mlx-nn-init-normal"]], "mlx.nn.init.uniform": [[319, "mlx-nn-init-uniform"]], "mlx.nn.elu": [[320, "mlx-nn-elu"]], "mlx.nn.gelu": [[321, "mlx-nn-gelu"]], "mlx.nn.gelu_approx": [[322, "mlx-nn-gelu-approx"]], "mlx.nn.gelu_fast_approx": [[323, "mlx-nn-gelu-fast-approx"]], "mlx.nn.glu": [[324, "mlx-nn-glu"]], "mlx.nn.hardswish": [[325, "mlx-nn-hardswish"]], "mlx.nn.leaky_relu": [[326, "mlx-nn-leaky-relu"]], "mlx.nn.log_sigmoid": [[327, "mlx-nn-log-sigmoid"]], "mlx.nn.log_softmax": [[328, "mlx-nn-log-softmax"]], "mlx.nn.losses.binary_cross_entropy": [[329, "mlx-nn-losses-binary-cross-entropy"]], "mlx.nn.losses.cosine_similarity_loss": [[330, "mlx-nn-losses-cosine-similarity-loss"]], "mlx.nn.losses.cross_entropy": [[331, "mlx-nn-losses-cross-entropy"]], "mlx.nn.losses.gaussian_nll_loss": [[332, "mlx-nn-losses-gaussian-nll-loss"]], "mlx.nn.losses.hinge_loss": [[333, "mlx-nn-losses-hinge-loss"]], "mlx.nn.losses.huber_loss": [[334, "mlx-nn-losses-huber-loss"]], "mlx.nn.losses.kl_div_loss": [[335, "mlx-nn-losses-kl-div-loss"]], "mlx.nn.losses.l1_loss": [[336, "mlx-nn-losses-l1-loss"]], "mlx.nn.losses.log_cosh_loss": [[337, "mlx-nn-losses-log-cosh-loss"]], "mlx.nn.losses.margin_ranking_loss": [[338, "mlx-nn-losses-margin-ranking-loss"]], "mlx.nn.losses.mse_loss": [[339, "mlx-nn-losses-mse-loss"]], "mlx.nn.losses.nll_loss": [[340, "mlx-nn-losses-nll-loss"]], "mlx.nn.losses.smooth_l1_loss": [[341, "mlx-nn-losses-smooth-l1-loss"]], "mlx.nn.losses.triplet_loss": [[342, "mlx-nn-losses-triplet-loss"]], "mlx.nn.mish": [[343, "mlx-nn-mish"]], "mlx.nn.prelu": [[344, "mlx-nn-prelu"]], "mlx.nn.relu": [[345, "mlx-nn-relu"]], "mlx.nn.relu6": [[346, "mlx-nn-relu6"]], "mlx.nn.selu": [[347, "mlx-nn-selu"]], "mlx.nn.sigmoid": [[348, "mlx-nn-sigmoid"]], "mlx.nn.silu": [[349, "mlx-nn-silu"]], "mlx.nn.softmax": [[350, "mlx-nn-softmax"]], "mlx.nn.softplus": [[351, "mlx-nn-softplus"]], "mlx.nn.softshrink": [[352, "mlx-nn-softshrink"]], "mlx.nn.step": [[353, "mlx-nn-step"]], "mlx.nn.tanh": [[354, "mlx-nn-tanh"]], "Functions": [[355, "functions"]], "Initializers": [[356, "initializers"]], "Layers": [[357, "layers"]], "Loss Functions": [[358, "loss-functions"]], "Module": [[359, "module"]], "Optimizers": [[361, "optimizers"]], "mlx.optimizers.AdaDelta": [[362, "mlx-optimizers-adadelta"]], "mlx.optimizers.Adafactor": [[363, "mlx-optimizers-adafactor"]], "mlx.optimizers.Adagrad": [[364, "mlx-optimizers-adagrad"]], "mlx.optimizers.Adam": [[365, "mlx-optimizers-adam"]], "mlx.optimizers.AdamW": [[366, "mlx-optimizers-adamw"]], "mlx.optimizers.Adamax": [[367, "mlx-optimizers-adamax"]], "mlx.optimizers.Lion": [[368, "mlx-optimizers-lion"]], "mlx.optimizers.Optimizer.apply_gradients": [[369, "mlx-optimizers-optimizer-apply-gradients"]], "mlx.optimizers.Optimizer.init": [[370, "mlx-optimizers-optimizer-init"]], "mlx.optimizers.Optimizer.state": [[371, "mlx-optimizers-optimizer-state"]], "mlx.optimizers.Optimizer.update": [[372, "mlx-optimizers-optimizer-update"]], "mlx.optimizers.RMSprop": [[373, "mlx-optimizers-rmsprop"]], "mlx.optimizers.SGD": [[374, "mlx-optimizers-sgd"]], "mlx.optimizers.cosine_decay": [[375, "mlx-optimizers-cosine-decay"]], "mlx.optimizers.exponential_decay": [[376, "mlx-optimizers-exponential-decay"]], "mlx.optimizers.join_schedules": [[377, "mlx-optimizers-join-schedules"]], "mlx.optimizers.linear_schedule": [[378, "mlx-optimizers-linear-schedule"]], "mlx.optimizers.step_decay": [[379, "mlx-optimizers-step-decay"]], "Common Optimizers": [[380, "common-optimizers"]], "Optimizer": [[381, "optimizer"]], "Schedulers": [[382, "schedulers"]], "Random": [[383, "random"]], "Transforms": [[384, "transforms"]], "Tree Utils": [[385, "tree-utils"]], "Compilation": [[386, "compilation"]], "Basics of Compile": [[386, "basics-of-compile"]], "Example Speedup": [[386, "example-speedup"]], "Debugging": [[386, "debugging"]], "Pure Functions": [[386, "pure-functions"]], "Compiling Training Graphs": [[386, "compiling-training-graphs"]], "Transformations with Compile": [[386, "transformations-with-compile"]], "Function Transforms": [[387, "function-transforms"]], "Automatic Differentiation": [[387, "automatic-differentiation"]], "Automatic Vectorization": [[387, "automatic-vectorization"]], "Indexing Arrays": [[388, "indexing-arrays"]], "Differences from NumPy": [[388, "differences-from-numpy"]], "In Place Updates": [[388, "in-place-updates"]], "Lazy Evaluation": [[389, "lazy-evaluation"]], "Why Lazy Evaluation": [[389, "why-lazy-evaluation"]], "Transforming Compute Graphs": [[389, "transforming-compute-graphs"]], "Only Compute What You Use": [[389, "only-compute-what-you-use"]], "When to Evaluate": [[389, "when-to-evaluate"]], "Conversion to NumPy and Other Frameworks": [[390, "conversion-to-numpy-and-other-frameworks"]], "PyTorch": [[390, "pytorch"]], "JAX": [[390, "jax"]], "TensorFlow": [[390, "tensorflow"]], "Quick Start Guide": [[391, "quick-start-guide"]], "Basics": [[391, "basics"]], "Function and Graph Transformations": [[391, "function-and-graph-transformations"]], "Saving and Loading Arrays": [[392, "saving-and-loading-arrays"]], "Serialization Formats": [[392, "id1"]], "Unified Memory": [[393, "unified-memory"]], "A Simple Example": [[393, "a-simple-example"]], "Using Streams": [[394, "using-streams"]], "Specifying the Stream": [[394, "specifying-the-stream"]]}, "indexentries": {"device (class in mlx.core)": [[8, "mlx.core.Device"]], "__init__() (device method)": [[8, "mlx.core.Device.__init__"]], "dtype (class in mlx.core)": [[9, "mlx.core.Dtype"]], "__init__() (dtype method)": [[9, "mlx.core.Dtype.__init__"]], "dtypecategory (class in mlx.core)": [[10, "mlx.core.DtypeCategory"]], "__init__() (dtypecategory method)": [[10, "mlx.core.DtypeCategory.__init__"]], "abs() (in module mlx.core)": [[11, "mlx.core.abs"]], "add() (in module mlx.core)": [[12, "mlx.core.add"]], "all() (in module mlx.core)": [[13, "mlx.core.all"]], "allclose() (in module mlx.core)": [[14, "mlx.core.allclose"]], "any() (in module mlx.core)": [[15, "mlx.core.any"]], "arange() (in module mlx.core)": [[16, "mlx.core.arange"]], "arccos() (in module mlx.core)": [[17, "mlx.core.arccos"]], "arccosh() (in module mlx.core)": [[18, "mlx.core.arccosh"]], "arcsin() (in module mlx.core)": [[19, "mlx.core.arcsin"]], "arcsinh() (in module mlx.core)": [[20, "mlx.core.arcsinh"]], "arctan() (in module mlx.core)": [[21, "mlx.core.arctan"]], "arctanh() (in module mlx.core)": [[22, "mlx.core.arctanh"]], "argmax() (in module mlx.core)": [[23, "mlx.core.argmax"]], "argmin() (in module mlx.core)": [[24, "mlx.core.argmin"]], "argpartition() (in module mlx.core)": [[25, "mlx.core.argpartition"]], "argsort() (in module mlx.core)": [[26, "mlx.core.argsort"]], "__init__() (array method)": [[27, "mlx.core.array.__init__"]], "array (class in mlx.core)": [[27, "mlx.core.array"]], "t (array property)": [[28, "mlx.core.array.T"]], "abs() (array method)": [[29, "mlx.core.array.abs"]], "all() (array method)": [[30, "mlx.core.array.all"]], "any() (array method)": [[31, "mlx.core.array.any"]], "argmax() (array method)": [[32, "mlx.core.array.argmax"]], "argmin() (array method)": [[33, "mlx.core.array.argmin"]], "astype() (array method)": [[34, "mlx.core.array.astype"]], "at (array property)": [[35, "mlx.core.array.at"]], "cos() (array method)": [[36, "mlx.core.array.cos"]], "cummax() (array method)": [[37, "mlx.core.array.cummax"]], "cummin() (array method)": [[38, "mlx.core.array.cummin"]], "cumprod() (array method)": [[39, "mlx.core.array.cumprod"]], "cumsum() (array method)": [[40, "mlx.core.array.cumsum"]], "diag() (array method)": [[41, "mlx.core.array.diag"]], "diagonal() (array method)": [[42, "mlx.core.array.diagonal"]], "dtype (array property)": [[43, "mlx.core.array.dtype"]], "exp() (array method)": [[44, "mlx.core.array.exp"]], "flatten() (array method)": [[45, "mlx.core.array.flatten"]], "item() (array method)": [[46, "mlx.core.array.item"]], "itemsize (array property)": [[47, "mlx.core.array.itemsize"]], "log() (array method)": [[48, "mlx.core.array.log"]], "log10() (array method)": [[49, "mlx.core.array.log10"]], "log1p() (array method)": [[50, "mlx.core.array.log1p"]], "log2() (array method)": [[51, "mlx.core.array.log2"]], "logsumexp() (array method)": [[52, "mlx.core.array.logsumexp"]], "max() (array method)": [[53, "mlx.core.array.max"]], "mean() (array method)": [[54, "mlx.core.array.mean"]], "min() (array method)": [[55, "mlx.core.array.min"]], "moveaxis() (array method)": [[56, "mlx.core.array.moveaxis"]], "nbytes (array property)": [[57, "mlx.core.array.nbytes"]], "ndim (array property)": [[58, "mlx.core.array.ndim"]], "prod() (array method)": [[59, "mlx.core.array.prod"]], "reciprocal() (array method)": [[60, "mlx.core.array.reciprocal"]], "reshape() (array method)": [[61, "mlx.core.array.reshape"]], "round() (array method)": [[62, "mlx.core.array.round"]], "rsqrt() (array method)": [[63, "mlx.core.array.rsqrt"]], "shape (array property)": [[64, "mlx.core.array.shape"]], "sin() (array method)": [[65, "mlx.core.array.sin"]], "size (array property)": [[66, "mlx.core.array.size"]], "split() (array method)": [[67, "mlx.core.array.split"]], "sqrt() (array method)": [[68, "mlx.core.array.sqrt"]], "square() (array method)": [[69, "mlx.core.array.square"]], "squeeze() (array method)": [[70, "mlx.core.array.squeeze"]], "sum() (array method)": [[71, "mlx.core.array.sum"]], "swapaxes() (array method)": [[72, "mlx.core.array.swapaxes"]], "tolist() (array method)": [[73, "mlx.core.array.tolist"]], "transpose() (array method)": [[74, "mlx.core.array.transpose"]], "var() (array method)": [[75, "mlx.core.array.var"]], "array_equal() (in module mlx.core)": [[76, "mlx.core.array_equal"]], "atleast_1d() (in module mlx.core)": [[77, "mlx.core.atleast_1d"]], "atleast_2d() (in module mlx.core)": [[78, "mlx.core.atleast_2d"]], "atleast_3d() (in module mlx.core)": [[79, "mlx.core.atleast_3d"]], "broadcast_to() (in module mlx.core)": [[80, "mlx.core.broadcast_to"]], "ceil() (in module mlx.core)": [[81, "mlx.core.ceil"]], "clip() (in module mlx.core)": [[82, "mlx.core.clip"]], "compile() (in module mlx.core)": [[83, "mlx.core.compile"]], "concatenate() (in module mlx.core)": [[84, "mlx.core.concatenate"]], "conv1d() (in module mlx.core)": [[85, "mlx.core.conv1d"]], "conv2d() (in module mlx.core)": [[86, "mlx.core.conv2d"]], "conv_general() (in module mlx.core)": [[87, "mlx.core.conv_general"]], "convolve() (in module mlx.core)": [[88, "mlx.core.convolve"]], "cos() (in module mlx.core)": [[89, "mlx.core.cos"]], "cosh() (in module mlx.core)": [[90, "mlx.core.cosh"]], "cummax() (in module mlx.core)": [[91, "mlx.core.cummax"]], "cummin() (in module mlx.core)": [[92, "mlx.core.cummin"]], "cumprod() (in module mlx.core)": [[93, "mlx.core.cumprod"]], "cumsum() (in module mlx.core)": [[94, "mlx.core.cumsum"]], "default_device() (in module mlx.core)": [[95, "mlx.core.default_device"]], "default_stream() (in module mlx.core)": [[96, "mlx.core.default_stream"]], "dequantize() (in module mlx.core)": [[97, "mlx.core.dequantize"]], "diag() (in module mlx.core)": [[98, "mlx.core.diag"]], "diagonal() (in module mlx.core)": [[99, "mlx.core.diagonal"]], "disable_compile() (in module mlx.core)": [[100, "mlx.core.disable_compile"]], "divide() (in module mlx.core)": [[101, "mlx.core.divide"]], "divmod() (in module mlx.core)": [[102, "mlx.core.divmod"]], "enable_compile() (in module mlx.core)": [[103, "mlx.core.enable_compile"]], "equal() (in module mlx.core)": [[104, "mlx.core.equal"]], "erf() (in module mlx.core)": [[105, "mlx.core.erf"]], "erfinv() (in module mlx.core)": [[106, "mlx.core.erfinv"]], "eval() (in module mlx.core)": [[107, "mlx.core.eval"]], "exp() (in module mlx.core)": [[108, "mlx.core.exp"]], "expand_dims() (in module mlx.core)": [[109, "mlx.core.expand_dims"]], "expm1() (in module mlx.core)": [[110, "mlx.core.expm1"]], "eye() (in module mlx.core)": [[111, "mlx.core.eye"]], "layer_norm() (in module mlx.core.fast)": [[112, "mlx.core.fast.layer_norm"]], "rms_norm() (in module mlx.core.fast)": [[113, "mlx.core.fast.rms_norm"]], "rope() (in module mlx.core.fast)": [[114, "mlx.core.fast.rope"]], "scaled_dot_product_attention() (in module mlx.core.fast)": [[115, "mlx.core.fast.scaled_dot_product_attention"]], "fft() (in module mlx.core.fft)": [[116, "mlx.core.fft.fft"]], "fft2() (in module mlx.core.fft)": [[117, "mlx.core.fft.fft2"]], "fftn() (in module mlx.core.fft)": [[118, "mlx.core.fft.fftn"]], "ifft() (in module mlx.core.fft)": [[119, "mlx.core.fft.ifft"]], "ifft2() (in module mlx.core.fft)": [[120, "mlx.core.fft.ifft2"]], "ifftn() (in module mlx.core.fft)": [[121, "mlx.core.fft.ifftn"]], "irfft() (in module mlx.core.fft)": [[122, "mlx.core.fft.irfft"]], "irfft2() (in module mlx.core.fft)": [[123, "mlx.core.fft.irfft2"]], "irfftn() (in module mlx.core.fft)": [[124, "mlx.core.fft.irfftn"]], "rfft() (in module mlx.core.fft)": [[125, "mlx.core.fft.rfft"]], "rfft2() (in module mlx.core.fft)": [[126, "mlx.core.fft.rfft2"]], "rfftn() (in module mlx.core.fft)": [[127, "mlx.core.fft.rfftn"]], "flatten() (in module mlx.core)": [[128, "mlx.core.flatten"]], "floor() (in module mlx.core)": [[129, "mlx.core.floor"]], "floor_divide() (in module mlx.core)": [[130, "mlx.core.floor_divide"]], "full() (in module mlx.core)": [[131, "mlx.core.full"]], "grad() (in module mlx.core)": [[132, "mlx.core.grad"]], "greater() (in module mlx.core)": [[133, "mlx.core.greater"]], "greater_equal() (in module mlx.core)": [[134, "mlx.core.greater_equal"]], "identity() (in module mlx.core)": [[135, "mlx.core.identity"]], "inner() (in module mlx.core)": [[136, "mlx.core.inner"]], "isclose() (in module mlx.core)": [[137, "mlx.core.isclose"]], "isinf() (in module mlx.core)": [[138, "mlx.core.isinf"]], "isnan() (in module mlx.core)": [[139, "mlx.core.isnan"]], "isneginf() (in module mlx.core)": [[140, "mlx.core.isneginf"]], "isposinf() (in module mlx.core)": [[141, "mlx.core.isposinf"]], "issubdtype() (in module mlx.core)": [[142, "mlx.core.issubdtype"]], "jvp() (in module mlx.core)": [[143, "mlx.core.jvp"]], "less() (in module mlx.core)": [[144, "mlx.core.less"]], "less_equal() (in module mlx.core)": [[145, "mlx.core.less_equal"]], "norm() (in module mlx.core.linalg)": [[146, "mlx.core.linalg.norm"]], "qr() (in module mlx.core.linalg)": [[147, "mlx.core.linalg.qr"]], "linspace() (in module mlx.core)": [[148, "mlx.core.linspace"]], "load() (in module mlx.core)": [[149, "mlx.core.load"]], "log() (in module mlx.core)": [[150, "mlx.core.log"]], "log10() (in module mlx.core)": [[151, "mlx.core.log10"]], "log1p() (in module mlx.core)": [[152, "mlx.core.log1p"]], "log2() (in module mlx.core)": [[153, "mlx.core.log2"]], "logaddexp() (in module mlx.core)": [[154, "mlx.core.logaddexp"]], "logical_and() (in module mlx.core)": [[155, "mlx.core.logical_and"]], "logical_not() (in module mlx.core)": [[156, "mlx.core.logical_not"]], "logical_or() (in module mlx.core)": [[157, "mlx.core.logical_or"]], "logsumexp() (in module mlx.core)": [[158, "mlx.core.logsumexp"]], "matmul() (in module mlx.core)": [[159, "mlx.core.matmul"]], "max() (in module mlx.core)": [[160, "mlx.core.max"]], "maximum() (in module mlx.core)": [[161, "mlx.core.maximum"]], "mean() (in module mlx.core)": [[162, "mlx.core.mean"]], "meshgrid() (in module mlx.core)": [[163, "mlx.core.meshgrid"]], "get_active_memory() (in module mlx.core.metal)": [[164, "mlx.core.metal.get_active_memory"]], "get_cache_memory() (in module mlx.core.metal)": [[165, "mlx.core.metal.get_cache_memory"]], "get_peak_memory() (in module mlx.core.metal)": [[166, "mlx.core.metal.get_peak_memory"]], "is_available() (in module mlx.core.metal)": [[167, "mlx.core.metal.is_available"]], "set_cache_limit() (in module mlx.core.metal)": [[168, "mlx.core.metal.set_cache_limit"]], "set_memory_limit() (in module mlx.core.metal)": [[169, "mlx.core.metal.set_memory_limit"]], "start_capture() (in module mlx.core.metal)": [[170, "mlx.core.metal.start_capture"]], "stop_capture() (in module mlx.core.metal)": [[171, "mlx.core.metal.stop_capture"]], "min() (in module mlx.core)": [[172, "mlx.core.min"]], "minimum() (in module mlx.core)": [[173, "mlx.core.minimum"]], "moveaxis() (in module mlx.core)": [[174, "mlx.core.moveaxis"]], "multiply() (in module mlx.core)": [[175, "mlx.core.multiply"]], "negative() (in module mlx.core)": [[176, "mlx.core.negative"]], "new_stream() (in module mlx.core)": [[177, "mlx.core.new_stream"]], "ones() (in module mlx.core)": [[178, "mlx.core.ones"]], "ones_like() (in module mlx.core)": [[179, "mlx.core.ones_like"]], "outer() (in module mlx.core)": [[180, "mlx.core.outer"]], "pad() (in module mlx.core)": [[181, "mlx.core.pad"]], "partition() (in module mlx.core)": [[182, "mlx.core.partition"]], "prod() (in module mlx.core)": [[183, "mlx.core.prod"]], "quantize() (in module mlx.core)": [[184, "mlx.core.quantize"]], "quantized_matmul() (in module mlx.core)": [[185, "mlx.core.quantized_matmul"]], "bernoulli() (in module mlx.core.random)": [[186, "mlx.core.random.bernoulli"]], "categorical() (in module mlx.core.random)": [[187, "mlx.core.random.categorical"]], "gumbel() (in module mlx.core.random)": [[188, "mlx.core.random.gumbel"]], "key() (in module mlx.core.random)": [[189, "mlx.core.random.key"]], "multivariate_normal() (in module mlx.core.random)": [[190, "mlx.core.random.multivariate_normal"]], "normal() (in module mlx.core.random)": [[191, "mlx.core.random.normal"]], "randint() (in module mlx.core.random)": [[192, "mlx.core.random.randint"]], "seed() (in module mlx.core.random)": [[193, "mlx.core.random.seed"]], "split() (in module mlx.core.random)": [[194, "mlx.core.random.split"]], "truncated_normal() (in module mlx.core.random)": [[195, "mlx.core.random.truncated_normal"]], "uniform() (in module mlx.core.random)": [[196, "mlx.core.random.uniform"]], "reciprocal() (in module mlx.core)": [[197, "mlx.core.reciprocal"]], "repeat() (in module mlx.core)": [[198, "mlx.core.repeat"]], "reshape() (in module mlx.core)": [[199, "mlx.core.reshape"]], "round() (in module mlx.core)": [[200, "mlx.core.round"]], "rsqrt() (in module mlx.core)": [[201, "mlx.core.rsqrt"]], "save() (in module mlx.core)": [[202, "mlx.core.save"]], "save_gguf() (in module mlx.core)": [[203, "mlx.core.save_gguf"]], "save_safetensors() (in module mlx.core)": [[204, "mlx.core.save_safetensors"]], "savez() (in module mlx.core)": [[205, "mlx.core.savez"]], "savez_compressed() (in module mlx.core)": [[206, "mlx.core.savez_compressed"]], "set_default_device() (in module mlx.core)": [[207, "mlx.core.set_default_device"]], "set_default_stream() (in module mlx.core)": [[208, "mlx.core.set_default_stream"]], "sigmoid() (in module mlx.core)": [[209, "mlx.core.sigmoid"]], "sign() (in module mlx.core)": [[210, "mlx.core.sign"]], "sin() (in module mlx.core)": [[211, "mlx.core.sin"]], "sinh() (in module mlx.core)": [[212, "mlx.core.sinh"]], "softmax() (in module mlx.core)": [[213, "mlx.core.softmax"]], "sort() (in module mlx.core)": [[214, "mlx.core.sort"]], "split() (in module mlx.core)": [[215, "mlx.core.split"]], "sqrt() (in module mlx.core)": [[216, "mlx.core.sqrt"]], "square() (in module mlx.core)": [[217, "mlx.core.square"]], "squeeze() (in module mlx.core)": [[218, "mlx.core.squeeze"]], "stack() (in module mlx.core)": [[219, "mlx.core.stack"]], "std() (in module mlx.core)": [[220, "mlx.core.std"]], "stop_gradient() (in module mlx.core)": [[221, "mlx.core.stop_gradient"]], "stream() (in module mlx.core)": [[222, "mlx.core.stream"]], "subtract() (in module mlx.core)": [[223, "mlx.core.subtract"]], "sum() (in module mlx.core)": [[224, "mlx.core.sum"]], "swapaxes() (in module mlx.core)": [[225, "mlx.core.swapaxes"]], "take() (in module mlx.core)": [[226, "mlx.core.take"]], "take_along_axis() (in module mlx.core)": [[227, "mlx.core.take_along_axis"]], "tan() (in module mlx.core)": [[228, "mlx.core.tan"]], "tanh() (in module mlx.core)": [[229, "mlx.core.tanh"]], "tensordot() (in module mlx.core)": [[230, "mlx.core.tensordot"]], "tile() (in module mlx.core)": [[231, "mlx.core.tile"]], "topk() (in module mlx.core)": [[232, "mlx.core.topk"]], "transpose() (in module mlx.core)": [[233, "mlx.core.transpose"]], "tri() (in module mlx.core)": [[234, "mlx.core.tri"]], "tril() (in module mlx.core)": [[235, "mlx.core.tril"]], "triu() (in module mlx.core)": [[236, "mlx.core.triu"]], "value_and_grad() (in module mlx.core)": [[237, "mlx.core.value_and_grad"]], "var() (in module mlx.core)": [[238, "mlx.core.var"]], "vjp() (in module mlx.core)": [[239, "mlx.core.vjp"]], "vmap() (in module mlx.core)": [[240, "mlx.core.vmap"]], "where() (in module mlx.core)": [[241, "mlx.core.where"]], "zeros() (in module mlx.core)": [[242, "mlx.core.zeros"]], "zeros_like() (in module mlx.core)": [[243, "mlx.core.zeros_like"]], "value_and_grad() (in module mlx.nn)": [[244, "mlx.nn.value_and_grad"]], "tree_flatten() (in module mlx.utils)": [[245, "mlx.utils.tree_flatten"]], "tree_map() (in module mlx.utils)": [[246, "mlx.utils.tree_map"]], "tree_unflatten() (in module mlx.utils)": [[247, "mlx.utils.tree_unflatten"]], "stream (class in mlx.core)": [[248, "mlx.core.Stream"]], "__init__() (stream method)": [[248, "mlx.core.Stream.__init__"]], "alibi (class in mlx.nn)": [[257, "mlx.nn.ALiBi"]], "avgpool1d (class in mlx.nn)": [[258, "mlx.nn.AvgPool1d"]], "avgpool2d (class in mlx.nn)": [[259, "mlx.nn.AvgPool2d"]], "batchnorm (class in mlx.nn)": [[260, "mlx.nn.BatchNorm"]], "conv1d (class in mlx.nn)": [[261, "mlx.nn.Conv1d"]], "conv2d (class in mlx.nn)": [[262, "mlx.nn.Conv2d"]], "dropout (class in mlx.nn)": [[263, "mlx.nn.Dropout"]], "dropout2d (class in mlx.nn)": [[264, "mlx.nn.Dropout2d"]], "dropout3d (class in mlx.nn)": [[265, "mlx.nn.Dropout3d"]], "embedding (class in mlx.nn)": [[266, "mlx.nn.Embedding"]], "gelu (class in mlx.nn)": [[267, "mlx.nn.GELU"]], "gru (class in mlx.nn)": [[268, "mlx.nn.GRU"]], "groupnorm (class in mlx.nn)": [[269, "mlx.nn.GroupNorm"]], "instancenorm (class in mlx.nn)": [[270, "mlx.nn.InstanceNorm"]], "lstm (class in mlx.nn)": [[271, "mlx.nn.LSTM"]], "layernorm (class in mlx.nn)": [[272, "mlx.nn.LayerNorm"]], "linear (class in mlx.nn)": [[273, "mlx.nn.Linear"]], "maxpool1d (class in mlx.nn)": [[274, "mlx.nn.MaxPool1d"]], "maxpool2d (class in mlx.nn)": [[275, "mlx.nn.MaxPool2d"]], "mish (class in mlx.nn)": [[276, "mlx.nn.Mish"]], "apply() (module method)": [[277, "mlx.nn.Module.apply"]], "apply_to_modules() (module method)": [[278, "mlx.nn.Module.apply_to_modules"]], "children() (module method)": [[279, "mlx.nn.Module.children"]], "eval() (module method)": [[280, "mlx.nn.Module.eval"]], "filter_and_map() (module method)": [[281, "mlx.nn.Module.filter_and_map"]], "freeze() (module method)": [[282, "mlx.nn.Module.freeze"]], "leaf_modules() (module method)": [[283, "mlx.nn.Module.leaf_modules"]], "load_weights() (module method)": [[284, "mlx.nn.Module.load_weights"]], "modules() (module method)": [[285, "mlx.nn.Module.modules"]], "named_modules() (module method)": [[286, "mlx.nn.Module.named_modules"]], "parameters() (module method)": [[287, "mlx.nn.Module.parameters"]], "save_weights() (module method)": [[288, "mlx.nn.Module.save_weights"]], "set_dtype() (module method)": [[289, "mlx.nn.Module.set_dtype"]], "state (module property)": [[290, "mlx.nn.Module.state"]], "train() (module method)": [[291, "mlx.nn.Module.train"]], "trainable_parameters() (module method)": [[292, "mlx.nn.Module.trainable_parameters"]], "training (module property)": [[293, "mlx.nn.Module.training"]], "unfreeze() (module method)": [[294, "mlx.nn.Module.unfreeze"]], "update() (module method)": [[295, "mlx.nn.Module.update"]], "update_modules() (module method)": [[296, "mlx.nn.Module.update_modules"]], "multiheadattention (class in mlx.nn)": [[297, "mlx.nn.MultiHeadAttention"]], "prelu (class in mlx.nn)": [[298, "mlx.nn.PReLU"]], "quantizedlinear (class in mlx.nn)": [[299, "mlx.nn.QuantizedLinear"]], "rmsnorm (class in mlx.nn)": [[300, "mlx.nn.RMSNorm"]], "rnn (class in mlx.nn)": [[301, "mlx.nn.RNN"]], "relu (class in mlx.nn)": [[302, "mlx.nn.ReLU"]], "rope (class in mlx.nn)": [[303, "mlx.nn.RoPE"]], "selu (class in mlx.nn)": [[304, "mlx.nn.SELU"]], "sequential (class in mlx.nn)": [[305, "mlx.nn.Sequential"]], "silu (class in mlx.nn)": [[306, "mlx.nn.SiLU"]], "sinusoidalpositionalencoding (class in mlx.nn)": [[307, "mlx.nn.SinusoidalPositionalEncoding"]], "softshrink (class in mlx.nn)": [[308, "mlx.nn.Softshrink"]], "step (class in mlx.nn)": [[309, "mlx.nn.Step"]], "transformer (class in mlx.nn)": [[310, "mlx.nn.Transformer"]], "upsample (class in mlx.nn)": [[311, "mlx.nn.Upsample"]], "constant() (in module mlx.nn.init)": [[312, "mlx.nn.init.constant"]], "glorot_normal() (in module mlx.nn.init)": [[313, "mlx.nn.init.glorot_normal"]], "glorot_uniform() (in module mlx.nn.init)": [[314, "mlx.nn.init.glorot_uniform"]], "he_normal() (in module mlx.nn.init)": [[315, "mlx.nn.init.he_normal"]], "he_uniform() (in module mlx.nn.init)": [[316, "mlx.nn.init.he_uniform"]], "identity() (in module mlx.nn.init)": [[317, "mlx.nn.init.identity"]], "normal() (in module mlx.nn.init)": [[318, "mlx.nn.init.normal"]], "uniform() (in module mlx.nn.init)": [[319, "mlx.nn.init.uniform"]], "elu() (in module mlx.nn)": [[320, "mlx.nn.elu"]], "gelu() (in module mlx.nn)": [[321, "mlx.nn.gelu"]], "gelu_approx() (in module mlx.nn)": [[322, "mlx.nn.gelu_approx"]], "gelu_fast_approx() (in module mlx.nn)": [[323, "mlx.nn.gelu_fast_approx"]], "glu() (in module mlx.nn)": [[324, "mlx.nn.glu"]], "hardswish() (in module mlx.nn)": [[325, "mlx.nn.hardswish"]], "leaky_relu() (in module mlx.nn)": [[326, "mlx.nn.leaky_relu"]], "log_sigmoid() (in module mlx.nn)": [[327, "mlx.nn.log_sigmoid"]], "log_softmax() (in module mlx.nn)": [[328, "mlx.nn.log_softmax"]], "binary_cross_entropy() (in module mlx.nn.losses)": [[329, "mlx.nn.losses.binary_cross_entropy"]], "cosine_similarity_loss() (in module mlx.nn.losses)": [[330, "mlx.nn.losses.cosine_similarity_loss"]], "cross_entropy() (in module mlx.nn.losses)": [[331, "mlx.nn.losses.cross_entropy"]], "gaussian_nll_loss() (in module mlx.nn.losses)": [[332, "mlx.nn.losses.gaussian_nll_loss"]], "hinge_loss() (in module mlx.nn.losses)": [[333, "mlx.nn.losses.hinge_loss"]], "huber_loss() (in module mlx.nn.losses)": [[334, "mlx.nn.losses.huber_loss"]], "kl_div_loss() (in module mlx.nn.losses)": [[335, "mlx.nn.losses.kl_div_loss"]], "l1_loss() (in module mlx.nn.losses)": [[336, "mlx.nn.losses.l1_loss"]], "log_cosh_loss() (in module mlx.nn.losses)": [[337, "mlx.nn.losses.log_cosh_loss"]], "margin_ranking_loss() (in module mlx.nn.losses)": [[338, "mlx.nn.losses.margin_ranking_loss"]], "mse_loss() (in module mlx.nn.losses)": [[339, "mlx.nn.losses.mse_loss"]], "nll_loss() (in module mlx.nn.losses)": [[340, "mlx.nn.losses.nll_loss"]], "smooth_l1_loss() (in module mlx.nn.losses)": [[341, "mlx.nn.losses.smooth_l1_loss"]], "triplet_loss() (in module mlx.nn.losses)": [[342, "mlx.nn.losses.triplet_loss"]], "mish() (in module mlx.nn)": [[343, "mlx.nn.mish"]], "prelu() (in module mlx.nn)": [[344, "mlx.nn.prelu"]], "relu() (in module mlx.nn)": [[345, "mlx.nn.relu"]], "relu6() (in module mlx.nn)": [[346, "mlx.nn.relu6"]], "selu() (in module mlx.nn)": [[347, "mlx.nn.selu"]], "sigmoid() (in module mlx.nn)": [[348, "mlx.nn.sigmoid"]], "silu() (in module mlx.nn)": [[349, "mlx.nn.silu"]], "softmax() (in module mlx.nn)": [[350, "mlx.nn.softmax"]], "softplus() (in module mlx.nn)": [[351, "mlx.nn.softplus"]], "softshrink() (in module mlx.nn)": [[352, "mlx.nn.softshrink"]], "step() (in module mlx.nn)": [[353, "mlx.nn.step"]], "tanh() (in module mlx.nn)": [[354, "mlx.nn.tanh"]], "module (class in mlx.nn)": [[359, "mlx.nn.Module"]], "adadelta (class in mlx.optimizers)": [[362, "mlx.optimizers.AdaDelta"]], "adafactor (class in mlx.optimizers)": [[363, "mlx.optimizers.Adafactor"]], "adagrad (class in mlx.optimizers)": [[364, "mlx.optimizers.Adagrad"]], "adam (class in mlx.optimizers)": [[365, "mlx.optimizers.Adam"]], "adamw (class in mlx.optimizers)": [[366, "mlx.optimizers.AdamW"]], "adamax (class in mlx.optimizers)": [[367, "mlx.optimizers.Adamax"]], "lion (class in mlx.optimizers)": [[368, "mlx.optimizers.Lion"]], "apply_gradients() (optimizer method)": [[369, "mlx.optimizers.Optimizer.apply_gradients"]], "init() (optimizer method)": [[370, "mlx.optimizers.Optimizer.init"]], "state (optimizer property)": [[371, "mlx.optimizers.Optimizer.state"]], "update() (optimizer method)": [[372, "mlx.optimizers.Optimizer.update"]], "rmsprop (class in mlx.optimizers)": [[373, "mlx.optimizers.RMSprop"]], "sgd (class in mlx.optimizers)": [[374, "mlx.optimizers.SGD"]], "cosine_decay() (in module mlx.optimizers)": [[375, "mlx.optimizers.cosine_decay"]], "exponential_decay() (in module mlx.optimizers)": [[376, "mlx.optimizers.exponential_decay"]], "join_schedules() (in module mlx.optimizers)": [[377, "mlx.optimizers.join_schedules"]], "linear_schedule() (in module mlx.optimizers)": [[378, "mlx.optimizers.linear_schedule"]], "step_decay() (in module mlx.optimizers)": [[379, "mlx.optimizers.step_decay"]], "optimizer (class in mlx.optimizers)": [[381, "mlx.optimizers.Optimizer"]]}})