mlx/docs/build/html/searchindex.js

Search.setIndex({"docnames": ["cpp/ops", "dev/extensions", "examples/linear_regression", "examples/llama-inference", "examples/mlp", "index", "install", "python/_autosummary/mlx.core.Device", "python/_autosummary/mlx.core.Dtype", "python/_autosummary/mlx.core.Stream", "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.cos", "python/_autosummary/mlx.core.array.dtype", "python/_autosummary/mlx.core.array.exp", "python/_autosummary/mlx.core.array.item", "python/_autosummary/mlx.core.array.log", "python/_autosummary/mlx.core.array.log1p", "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.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.sum", "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.broadcast_to", "python/_autosummary/mlx.core.ceil", "python/_autosummary/mlx.core.clip", "python/_autosummary/mlx.core.concatenate", "python/_autosummary/mlx.core.conv1d", "python/_autosummary/mlx.core.conv2d", "python/_autosummary/mlx.core.convolve", "python/_autosummary/mlx.core.cos", "python/_autosummary/mlx.core.cosh", "python/_autosummary/mlx.core.default_device", "python/_autosummary/mlx.core.default_stream", "python/_autosummary/mlx.core.dequantize", "python/_autosummary/mlx.core.divide", "python/_autosummary/mlx.core.divmod", "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.eye", "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.jvp", "python/_autosummary/mlx.core.less", "python/_autosummary/mlx.core.less_equal", "python/_autosummary/mlx.core.linalg.norm", "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.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.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.simplify", "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.stop_gradient", "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.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.optimizers.AdaDelta", "python/_autosummary/mlx.optimizers.Adagrad", "python/_autosummary/mlx.optimizers.Adam", "python/_autosummary/mlx.optimizers.AdamW", "python/_autosummary/mlx.optimizers.Adamax", "python/_autosummary/mlx.optimizers.Lion", "python/_autosummary/mlx.optimizers.Optimizer", "python/_autosummary/mlx.optimizers.OptimizerState", "python/_autosummary/mlx.optimizers.RMSprop", "python/_autosummary/mlx.optimizers.SGD", "python/_autosummary/mlx.utils.tree_flatten", "python/_autosummary/mlx.utils.tree_map", "python/_autosummary/mlx.utils.tree_unflatten", "python/array", "python/data_types", "python/devices_and_streams", "python/fft", "python/linalg", "python/nn", "python/nn/_autosummary/mlx.nn.ALiBi", "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.GroupNorm", "python/nn/_autosummary/mlx.nn.InstanceNorm", "python/nn/_autosummary/mlx.nn.LayerNorm", "python/nn/_autosummary/mlx.nn.Linear", "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.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.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.Step", "python/nn/_autosummary/mlx.nn.Transformer", "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.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.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.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.selu", "python/nn/_autosummary_functions/mlx.nn.silu", "python/nn/_autosummary_functions/mlx.nn.step", "python/nn/functions", "python/nn/layers", "python/nn/losses", "python/nn/module", "python/ops", "python/optimizers", "python/random", "python/transforms", "python/tree_utils", "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", "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.Stream.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.cos.rst", "python/_autosummary/mlx.core.array.dtype.rst", "python/_autosummary/mlx.core.array.exp.rst", "python/_autosummary/mlx.core.array.item.rst", "python/_autosummary/mlx.core.array.log.rst", "python/_autosummary/mlx.core.array.log1p.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.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.sum.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.broadcast_to.rst", "python/_autosummary/mlx.core.ceil.rst", "python/_autosummary/mlx.core.clip.rst", "python/_autosummary/mlx.core.concatenate.rst", "python/_autosummary/mlx.core.conv1d.rst", "python/_autosummary/mlx.core.conv2d.rst", "python/_autosummary/mlx.core.convolve.rst", "python/_autosummary/mlx.core.cos.rst", "python/_autosummary/mlx.core.cosh.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.divide.rst", "python/_autosummary/mlx.core.divmod.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.eye.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.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.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.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.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.simplify.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.stop_gradient.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.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.optimizers.AdaDelta.rst", "python/_autosummary/mlx.optimizers.Adagrad.rst", "python/_autosummary/mlx.optimizers.Adam.rst", "python/_autosummary/mlx.optimizers.AdamW.rst", "python/_autosummary/mlx.optimizers.Adamax.rst", "python/_autosummary/mlx.optimizers.Lion.rst", "python/_autosummary/mlx.optimizers.Optimizer.rst", "python/_autosummary/mlx.optimizers.OptimizerState.rst", "python/_autosummary/mlx.optimizers.RMSprop.rst", "python/_autosummary/mlx.optimizers.SGD.rst", "python/_autosummary/mlx.utils.tree_flatten.rst", "python/_autosummary/mlx.utils.tree_map.rst", "python/_autosummary/mlx.utils.tree_unflatten.rst", "python/array.rst", "python/data_types.rst", "python/devices_and_streams.rst", "python/fft.rst", "python/linalg.rst", "python/nn.rst", "python/nn/_autosummary/mlx.nn.ALiBi.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.GroupNorm.rst", "python/nn/_autosummary/mlx.nn.InstanceNorm.rst", "python/nn/_autosummary/mlx.nn.LayerNorm.rst", "python/nn/_autosummary/mlx.nn.Linear.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.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.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.Step.rst", "python/nn/_autosummary/mlx.nn.Transformer.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.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.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.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.selu.rst", "python/nn/_autosummary_functions/mlx.nn.silu.rst", "python/nn/_autosummary_functions/mlx.nn.step.rst", "python/nn/functions.rst", "python/nn/layers.rst", "python/nn/losses.rst", "python/nn/module.rst", "python/ops.rst", "python/optimizers.rst", "python/random.rst", "python/transforms.rst", "python/tree_utils.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", "Linear Regression", "LLM inference", "Multi-Layer Perceptron", "MLX", "Build and Install", "mlx.core.Device", "mlx.core.Dtype", "mlx.core.Stream", "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.cos", "mlx.core.array.dtype", "mlx.core.array.exp", "mlx.core.array.item", "mlx.core.array.log", "mlx.core.array.log1p", "mlx.core.array.logsumexp", "mlx.core.array.max", "mlx.core.array.mean", "mlx.core.array.min", "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.sum", "mlx.core.array.tolist", "mlx.core.array.transpose", "mlx.core.array.var", "mlx.core.array_equal", "mlx.core.broadcast_to", "mlx.core.ceil", "mlx.core.clip", "mlx.core.concatenate", "mlx.core.conv1d", "mlx.core.conv2d", "mlx.core.convolve", "mlx.core.cos", "mlx.core.cosh", "mlx.core.default_device", "mlx.core.default_stream", "mlx.core.dequantize", "mlx.core.divide", "mlx.core.divmod", "mlx.core.equal", "mlx.core.erf", "mlx.core.erfinv", "mlx.core.eval", "mlx.core.exp", "mlx.core.expand_dims", "mlx.core.eye", "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.jvp", "mlx.core.less", "mlx.core.less_equal", "mlx.core.linalg.norm", "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.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.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.simplify", "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.stop_gradient", "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.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.optimizers.AdaDelta", "mlx.optimizers.Adagrad", "mlx.optimizers.Adam", "mlx.optimizers.AdamW", "mlx.optimizers.Adamax", "mlx.optimizers.Lion", "mlx.optimizers.Optimizer", "mlx.optimizers.OptimizerState", "mlx.optimizers.RMSprop", "mlx.optimizers.SGD", "mlx.utils.tree_flatten", "mlx.utils.tree_map", "mlx.utils.tree_unflatten", "Array", "Data Types", "Devices and Streams", "FFT", "Linear Algebra", "Neural Networks", "mlx.nn.ALiBi", "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.GroupNorm", "mlx.nn.InstanceNorm", "mlx.nn.LayerNorm", "mlx.nn.Linear", "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.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.ReLU", "mlx.nn.RoPE", "mlx.nn.SELU", "mlx.nn.Sequential", "mlx.nn.SiLU", "mlx.nn.SinusoidalPositionalEncoding", "mlx.nn.Step", "mlx.nn.Transformer", "mlx.nn.gelu", "mlx.nn.gelu_approx", "mlx.nn.gelu_fast_approx", "mlx.nn.losses.binary_cross_entropy", "mlx.nn.losses.cosine_similarity_loss", "mlx.nn.losses.cross_entropy", "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.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.selu", "mlx.nn.silu", "mlx.nn.step", "Functions", "Layers", "Loss Functions", "Module", "Operations", "Optimizers", "Random", "Transforms", "Tree Utils", "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, 6, 209, 278, 280, 281, 283, 284, 285, 286, 287, 288, 289], "provid": [1, 3, 72, 98, 178, 183, 202, 209, 224, 229, 231, 239, 240, 241, 244, 253, 275, 278, 288, 290], "open": [1, 6, 15, 141, 145], "flexibl": [1, 5, 241], "which": [1, 3, 4, 5, 6, 15, 33, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 98, 103, 106, 108, 134, 137, 138, 147, 148, 151, 152, 153, 154, 155, 168, 169, 174, 183, 185, 186, 215, 216, 218, 224, 228, 247, 259, 261, 267, 281, 284, 285, 289, 290], "user": [1, 3, 209], "mai": [1, 106, 215, 284], "add": [1, 3, 80, 113, 131, 134, 212, 213, 289], "special": 1, "without": [1, 3, 5, 170, 242, 275, 283, 285, 286, 289], "much": [1, 3, 285], "hassl": 1, "while": [1, 3, 6, 148, 247, 285, 286], "librari": [1, 6, 209], "suppli": 1, "effici": [1, 3, 5, 215, 247, 285, 287], "can": [1, 3, 5, 6, 11, 15, 47, 58, 73, 74, 75, 78, 99, 100, 104, 105, 106, 113, 120, 123, 125, 136, 137, 141, 144, 145, 171, 183, 209, 217, 228, 239, 249, 278, 280, 281, 283, 284, 285, 286, 287, 288, 289, 290], "compos": [1, 5, 209, 287], "ani": [1, 3, 5, 15, 160, 201, 202, 203, 209, 218, 224, 225, 228, 244, 253, 283, 285, 287, 288, 289], "number": [1, 15, 52, 66, 72, 81, 98, 101, 103, 107, 131, 134, 135, 137, 140, 143, 145, 147, 149, 160, 178, 180, 183, 185, 186, 209, 211, 212, 213, 215, 216, 219, 220, 242, 243, 253, 281, 290], "applic": [1, 6], "aris": [1, 286], "case": [1, 3, 84, 87, 88, 90, 91, 92, 93, 94, 118, 148, 168, 215, 248, 252, 267, 272, 274, 287, 288, 289, 290], "where": [1, 4, 81, 134, 183, 186, 211, 212, 213, 214, 215, 216, 218, 219, 220, 221, 222, 228, 243, 245, 248, 250, 252, 254, 255, 256, 270, 272, 273, 274, 284], "new": [1, 4, 61, 124, 148, 169, 179, 202, 242, 278, 280, 284, 285, 286], "function": [1, 2, 3, 4, 5, 13, 74, 76, 77, 98, 103, 106, 118, 158, 160, 183, 185, 186, 190, 202, 209, 218, 223, 225, 229, 239, 243, 249, 252, 253, 254, 255, 256, 269, 274, 280, 281, 283, 284, 285, 286, 288], "highli": [1, 6], "optim": [1, 2, 4, 5, 240, 285], "ar": [1, 2, 3, 4, 5, 6, 13, 15, 60, 61, 63, 67, 81, 83, 84, 86, 87, 89, 90, 92, 93, 98, 103, 106, 108, 118, 130, 131, 132, 134, 135, 136, 137, 138, 141, 144, 145, 154, 155, 168, 169, 174, 183, 185, 186, 196, 201, 202, 211, 212, 213, 214, 215, 216, 219, 220, 221, 222, 231, 242, 244, 275, 278, 283, 284, 285, 286, 287, 288, 289], "need": [1, 3, 4, 5, 60, 134, 209, 240, 241, 251, 253, 281, 285, 286, 287, 289], "For": [1, 3, 6, 106, 134, 203, 209, 211, 215, 224, 229, 236, 239, 244, 247, 251, 281, 284, 285, 286, 287, 288, 289], "you": [1, 3, 4, 5, 6, 209, 251, 253, 281, 284, 286, 288, 289], "design": [1, 2, 5, 281, 289], "your": [1, 3, 6, 278, 285], "own": [1, 6, 286], "link": [1, 6], "top": [1, 222], "core": [1, 2, 3, 4, 209, 211, 220, 231, 234, 237, 257, 278, 280, 286, 287], "we": [1, 2, 3, 4, 72, 134, 135, 194, 196, 209, 217, 249, 281, 283, 285, 289], "inner": 1, "work": [1, 3, 6, 284, 285], "go": [1, 3], "over": [1, 3, 4, 12, 14, 22, 23, 24, 25, 65, 66, 84, 87, 90, 93, 102, 106, 107, 117, 119, 121, 122, 132, 133, 150, 163, 164, 172, 178, 184, 211, 212, 213, 219, 221, 245, 259], "simpl": [1, 3, 4, 209, 217, 275, 285], "learn": [1, 2, 4, 5, 191, 192, 193, 194, 195, 196, 199, 200, 211, 219, 220, 221, 243, 245], "step": [1, 3, 4, 15, 209], "involv": [1, 280], "ad": [1, 2, 6, 191, 192, 193, 194, 195, 199, 220, 278, 285, 288], "let": [1, 2, 3, 285, 286], "s": [1, 2, 3, 4, 35, 44, 72, 83, 84, 86, 87, 89, 90, 92, 93, 98, 106, 108, 121, 130, 134, 137, 149, 152, 153, 183, 184, 186, 190, 197, 209, 228, 229, 231, 235, 239, 280, 281, 285, 286, 287, 288, 289], "sai": [1, 3, 285], "would": [1, 3, 284, 285, 286, 289], "like": [1, 3, 5, 129, 189, 216, 264, 285, 286, 287, 289], "an": [1, 3, 4, 6, 8, 12, 14, 26, 61, 65, 66, 78, 81, 94, 97, 101, 106, 119, 122, 124, 128, 129, 131, 133, 134, 135, 147, 148, 149, 165, 168, 173, 174, 175, 178, 180, 186, 188, 189, 191, 197, 198, 201, 202, 209, 214, 219, 221, 222, 224, 242, 243, 244, 253, 255, 270, 281, 283, 284, 285, 286, 287, 288, 289, 290], "take": [1, 3, 4, 98, 103, 120, 123, 129, 135, 175, 183, 185, 186, 189, 242, 281, 284, 288, 289, 290], "two": [1, 11, 13, 60, 73, 75, 83, 86, 92, 99, 100, 104, 105, 113, 118, 120, 123, 125, 130, 173, 244, 258, 284, 289], "arrai": [1, 3, 4, 5, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 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, 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, 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, 209, 211, 224, 231, 234, 237, 243, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 274, 278, 285, 286, 287, 289], "x": [1, 2, 3, 4, 76, 101, 106, 135, 138, 149, 154, 158, 160, 181, 182, 187, 196, 202, 209, 211, 218, 219, 220, 221, 222, 223, 224, 243, 245, 246, 248, 250, 251, 252, 254, 255, 256, 267, 269, 270, 271, 272, 273, 274, 278, 280, 284, 285, 286, 287, 289], "y": [1, 2, 3, 4, 160, 187, 192, 209, 211, 215, 219, 220, 221, 222, 245, 260, 267, 280, 285, 286], "scale": [1, 3, 72, 134, 135, 215, 216, 242, 247, 248, 251, 272], "them": [1, 3, 209, 229, 239, 289], "both": [1, 11, 73, 74, 75, 99, 100, 104, 105, 106, 113, 120, 123, 125, 137, 171, 220, 280, 287, 289], "some": [1, 2, 3, 4, 229, 239, 285], "coeffici": [1, 191, 193, 194, 195, 196], "alpha": [1, 134, 194, 199, 248, 268, 270, 272], "beta": [1, 72, 134, 193, 194, 195, 196, 211, 219, 220, 221, 267], "respect": [1, 2, 4, 98, 134, 183, 202, 209, 211, 218, 219, 220, 221, 278, 287], "togeth": [1, 4, 134, 202], "get": [1, 2, 4, 6, 66, 139, 198, 209, 285, 289], "z": [1, 160, 285], "well": [1, 3, 209, 229, 239, 242, 285], "veri": [1, 3, 242, 285, 289], "easili": 1, "do": [1, 3, 6, 194, 209, 230, 239, 278, 285], "just": [1, 4, 284], "write": [1, 3, 209, 286], "out": [1, 6, 215, 216, 236, 284], "follow": [1, 3, 4, 5, 6, 15, 67, 72, 106, 134, 191, 192, 193, 194, 195, 196, 200, 209, 255, 256, 262, 281, 289], "import": [1, 2, 3, 4, 6, 106, 154, 160, 183, 201, 202, 203, 209, 211, 220, 231, 257, 278, 284, 285, 286, 287], "mx": [1, 2, 3, 4, 106, 108, 154, 160, 183, 209, 211, 220, 224, 231, 246, 257, 258, 262, 271, 278, 280, 281, 284, 285, 286, 287, 288, 289, 290], "def": [1, 2, 3, 4, 160, 183, 209, 278, 284, 285, 286, 289], "simple_axpbi": 1, "float": [1, 13, 15, 57, 96, 97, 106, 135, 136, 141, 144, 145, 191, 192, 193, 194, 195, 196, 199, 200, 205, 211, 214, 215, 216, 219, 220, 221, 224, 245, 247, 251, 252, 253, 258, 259, 261, 267, 268, 274], "return": [1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 37, 50, 57, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 72, 73, 74, 75, 76, 77, 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, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 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, 201, 202, 203, 209, 226, 228, 230, 232, 233, 234, 237, 244, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 278, 283, 284, 285, 286, 288, 289], "thi": [1, 3, 4, 6, 12, 13, 14, 15, 22, 23, 24, 25, 103, 106, 113, 117, 118, 119, 121, 122, 132, 133, 137, 160, 163, 164, 165, 172, 174, 184, 209, 214, 215, 216, 225, 226, 228, 229, 232, 233, 234, 237, 239, 240, 241, 242, 244, 252, 255, 256, 264, 274, 278, 283, 285, 286, 288], "perform": [1, 3, 5, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 118, 135, 149, 163, 174, 209, 219, 253, 284, 285, 289], "leav": [1, 202], "differenti": [1, 5], "howev": [1, 209, 218, 219, 281, 285, 286], "vector": [1, 2, 5, 102, 103, 106, 174, 185, 186, 217, 287], "math": [1, 3, 268], "often": [1, 216], "realiz": 1, "axpbi": 1, "routin": 1, "defin": [1, 2, 3, 4, 6, 106, 135, 198, 201, 286], "same": [1, 3, 6, 60, 61, 66, 67, 88, 91, 92, 93, 98, 103, 131, 137, 149, 185, 187, 209, 211, 214, 219, 220, 244, 268, 278, 281, 284, 289], "realli": 1, "part": [1, 284], "doe": [1, 3, 6, 209, 284, 285, 286], "fast": [1, 160, 218, 256, 289], "so": [1, 3, 6, 98, 160, 183, 214, 280, 285, 289], "decid": [1, 202, 228], "want": [1, 3, 289], "reli": 1, "acceler": [1, 211], "framework": [1, 5], "continu": 1, "impos": 1, "our": [1, 3, 4, 191, 192, 193, 195, 196, 249], "assumpt": 1, "also": [1, 3, 4, 5, 11, 73, 74, 75, 84, 87, 90, 93, 99, 100, 104, 105, 113, 120, 123, 125, 134, 171, 190, 198, 209, 228, 240, 242, 244, 248, 250, 254, 272, 273, 275, 280, 284, 285, 286, 287, 290], "assum": [1, 3, 202, 209, 219], "how": [1, 3, 4, 209, 212, 213, 217, 284, 289], "gradient": [1, 2, 4, 98, 170, 183, 190, 191, 193, 194, 195, 196, 200, 209, 229, 240, 244, 264, 278, 280, 284, 285, 286, 287], "ins": 1, "what": [1, 3, 202], "coincid": 1, "right": [1, 6, 134, 218, 255, 256, 261, 268], "place": [1, 3, 149, 285, 286], "cours": 1, "The": [1, 3, 4, 5, 6, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 35, 44, 50, 57, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 72, 73, 74, 75, 76, 77, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 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, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 143, 144, 145, 146, 147, 148, 152, 153, 158, 159, 161, 162, 163, 164, 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, 205, 211, 212, 213, 214, 215, 216, 217, 219, 220, 221, 222, 225, 231, 240, 241, 242, 244, 245, 247, 249, 251, 252, 253, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 274, 278, 280, 284, 285, 286, 287, 288, 289, 290], "structur": 1, "from": [1, 3, 4, 5, 72, 89, 90, 92, 93, 97, 106, 108, 118, 129, 134, 136, 137, 138, 139, 141, 144, 154, 168, 170, 171, 174, 175, 187, 189, 201, 202, 203, 209, 222, 229, 231, 242, 267, 283, 285, 286, 287, 288, 289], "frontend": 1, "api": 1, "redirect": 1, "when": [1, 3, 5, 6, 106, 108, 212, 213, 262, 267, 278, 281, 289], "appropri": 1, "fallback": 1, "metal": 1, "vjp": [1, 287], "jvp": [1, 287], "In": [1, 3, 4, 118, 134, 191, 192, 193, 195, 196, 202, 209, 215, 219, 278, 283, 285, 288, 289], "one": [1, 3, 6, 57, 63, 66, 80, 81, 106, 111, 118, 135, 137, 168, 171, 239, 289], "sentenc": 1, "comput": [1, 2, 3, 4, 5, 6, 72, 98, 103, 106, 113, 121, 130, 134, 160, 163, 170, 178, 183, 184, 185, 190, 191, 193, 194, 195, 196, 209, 211, 219, 220, 221, 229, 240, 244, 245, 247, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 280, 287, 289], "graph": [1, 3, 4, 5, 160], "rule": 1, "evalu": [1, 3, 4, 5, 78, 103, 185, 209, 227, 236, 278, 280, 287], "said": [1, 3], "start": [1, 2, 3, 5, 6, 15, 107, 165, 284, 289], "discuss": 1, "more": [1, 4, 8, 57, 118, 152, 153, 209, 211, 215, 247, 251, 281, 284, 287, 289], "detail": [1, 8, 191, 192, 193, 195, 196, 209, 215, 247, 251, 284, 287], "thei": [1, 2, 3, 67, 196, 249, 260, 278, 283, 285, 287, 288, 289], "c": [1, 3, 106, 205, 211, 212, 213, 215, 216, 220, 286, 287, 289], "scalar": [1, 11, 13, 26, 37, 57, 60, 61, 63, 73, 74, 75, 96, 97, 98, 99, 100, 104, 105, 106, 107, 113, 114, 115, 116, 118, 120, 123, 125, 131, 141, 144, 145, 171, 183, 187, 190, 268, 285, 287], "sum": [1, 2, 11, 102, 106, 117, 163, 178, 209, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 284, 286], "element": [1, 10, 11, 16, 17, 18, 19, 20, 21, 24, 52, 62, 68, 69, 72, 73, 74, 75, 76, 77, 79, 81, 95, 96, 99, 100, 104, 105, 109, 110, 111, 112, 113, 114, 115, 116, 120, 123, 125, 126, 132, 134, 135, 146, 147, 150, 158, 159, 161, 162, 166, 167, 171, 174, 176, 177, 183, 187, 214, 215, 216, 223, 243, 247, 250, 269, 270, 273], "wise": [1, 10, 11, 16, 17, 18, 19, 20, 21, 62, 68, 69, 73, 74, 75, 76, 77, 79, 95, 96, 99, 100, 104, 105, 109, 110, 111, 112, 113, 114, 115, 116, 120, 123, 125, 126, 146, 150, 158, 159, 161, 162, 166, 167, 171, 176, 177, 215, 216, 223, 243, 250, 269, 270, 273], "numpi": [1, 3, 4, 5, 11, 13, 15, 61, 73, 74, 75, 99, 100, 104, 105, 113, 118, 120, 123, 125, 171, 285, 287, 288], "style": [1, 11, 13, 73, 74, 75, 99, 100, 104, 105, 113, 118, 120, 123, 125, 171], "broadcast": [1, 11, 13, 61, 63, 73, 74, 75, 97, 99, 100, 104, 105, 113, 118, 120, 123, 125, 136, 137, 144, 145, 171, 175, 187, 242], "between": [1, 5, 63, 253, 258, 260, 261, 264, 285, 289], "input": [1, 2, 3, 10, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 73, 74, 75, 76, 77, 79, 80, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 98, 99, 100, 102, 103, 104, 105, 106, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 129, 130, 131, 132, 133, 134, 135, 143, 146, 147, 148, 149, 150, 158, 159, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 181, 182, 183, 184, 186, 187, 189, 211, 212, 213, 215, 216, 217, 219, 220, 221, 222, 242, 244, 245, 247, 252, 253, 257, 258, 260, 261, 262, 264, 266, 268, 274, 284, 287, 288], "upcast": 1, "const": 1, "factor": [1, 259], "streamordevic": 1, "stream": [1, 5, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 51, 53, 54, 55, 56, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 71, 72, 73, 74, 75, 76, 77, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 99, 100, 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, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 157, 158, 159, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 184, 187, 188, 189, 289], "schedul": [1, 289], "itself": 1, "call": [1, 3, 4, 27, 96, 209, 217, 229, 239, 249, 278, 280, 285], "other": [1, 3, 5, 106, 196, 209, 230, 278, 284, 285, 287], "within": [1, 24], "simplest": [1, 209], "wai": [1, 3, 6, 209, 284], "about": [1, 3, 4, 285, 289], "term": [1, 191, 192, 193, 194, 195, 199], "exist": [1, 3, 229, 239], "auto": [1, 6], "ax": [1, 12, 14, 22, 23, 58, 80, 83, 84, 86, 87, 89, 90, 92, 93, 102, 106, 117, 119, 121, 122, 131, 133, 163, 168, 172, 173, 178, 179, 184], "multipli": [1, 134, 135, 214, 251], "earlier": 1, "goal": 1, "themselv": 1, "contain": [1, 3, 50, 88, 89, 90, 106, 114, 115, 116, 134, 165, 187, 209, 228, 230, 231, 253, 278], "act": [1, 264], "data": [1, 4, 5, 8, 15, 81, 91, 92, 97, 101, 107, 128, 144, 180, 188, 216, 284, 286], "nor": [1, 98, 183], "rather": [1, 289], "easi": [1, 209], "interfac": 1, "block": [1, 3, 253], "A": [1, 3, 5, 6, 50, 60, 98, 103, 106, 108, 117, 118, 134, 136, 137, 138, 140, 141, 144, 145, 165, 169, 183, 185, 186, 190, 193, 195, 201, 202, 203, 209, 211, 215, 219, 220, 221, 223, 228, 232, 233, 240, 241, 245, 249, 251, 253, 256, 268, 269, 278, 280, 285, 286], "It": [1, 3, 6, 98, 183, 195, 197, 209, 241, 244, 286, 288], "creat": [1, 3, 6, 81, 101, 209, 278, 280, 284, 286], "output": [1, 3, 6, 12, 13, 14, 15, 22, 23, 24, 61, 81, 88, 91, 92, 93, 97, 98, 101, 106, 107, 117, 119, 121, 122, 128, 129, 132, 133, 136, 137, 138, 140, 141, 144, 145, 154, 155, 163, 168, 172, 175, 180, 183, 184, 185, 186, 187, 188, 189, 211, 212, 213, 220, 222, 242, 244, 252, 253, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 274, 284, 285, 286, 287, 288, 289], "given": [1, 12, 14, 24, 61, 63, 64, 72, 78, 80, 82, 83, 84, 85, 86, 87, 91, 92, 93, 97, 106, 117, 119, 121, 122, 133, 141, 149, 163, 165, 172, 180, 181, 182, 184, 214, 228, 242, 258], "set": [1, 3, 4, 6, 198, 218, 222, 227, 229, 236, 239, 240, 244, 247, 252, 258, 268, 274, 278, 281, 285], "further": [1, 6], "class": [1, 3, 4, 7, 8, 9, 26, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 278], "under": [1, 106], "These": [1, 175, 289], "word": 1, "bit": [1, 72, 134, 135, 205, 224, 244], "abstract": 1, "back": [1, 3, 286], "give": [1, 3, 4, 24], "ourselv": 1, "concret": [1, 222, 285, 289], "imag": [1, 213, 215, 216], "public": [1, 209], "explicit": [1, 281, 286], "alpha_": 1, "beta_": 1, "must": [1, 6, 63, 78, 97, 106, 136, 137, 141, 144, 145, 187, 286], "know": [1, 3], "popul": 1, "To": [1, 2, 3, 4, 6, 209, 287], "avoid": 1, "unnecessari": [1, 3], "alloc": [1, 278], "respons": 1, "space": [1, 107, 266], "void": 1, "eval_cpu": 1, "std": 1, "overrid": 1, "eval_gpu": 1, "jacobian": [1, 103, 185, 287], "product": [1, 102, 103, 118, 130, 133, 178, 185, 242, 287], "primal": [1, 103, 185], "tangent": [1, 20, 21, 103, 176, 177], "int": [1, 3, 4, 7, 9, 12, 14, 15, 22, 23, 24, 25, 29, 30, 31, 32, 40, 41, 42, 43, 45, 48, 50, 53, 56, 57, 59, 61, 64, 65, 66, 72, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 97, 98, 101, 106, 107, 117, 119, 121, 122, 124, 128, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 147, 148, 149, 163, 164, 165, 168, 169, 172, 173, 174, 175, 178, 179, 180, 181, 182, 183, 184, 186, 188, 209, 211, 212, 213, 217, 219, 220, 221, 222, 242, 244, 245, 247, 251, 253, 258, 259, 262, 266, 268, 278], "argnum": [1, 98, 183], "cotan": 1, "across": [1, 219], "pair": [1, 131, 231, 247], "repres": [1, 3, 268, 286], "axi": [1, 3, 4, 12, 14, 22, 23, 24, 25, 29, 30, 31, 32, 40, 41, 42, 43, 45, 53, 56, 59, 64, 80, 82, 85, 88, 89, 90, 91, 92, 93, 106, 117, 119, 121, 122, 124, 131, 132, 133, 137, 147, 163, 164, 165, 168, 169, 172, 173, 174, 175, 179, 184, 186, 258, 259, 262, 266, 268, 284], "correspond": [1, 12, 14, 57, 63, 72, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 117, 119, 122, 133, 172, 178, 186, 202], "dimens": [1, 3, 12, 14, 22, 23, 44, 50, 57, 66, 80, 89, 90, 92, 93, 94, 102, 106, 117, 118, 119, 121, 122, 133, 134, 137, 143, 172, 175, 178, 179, 184, 211, 212, 213, 215, 216, 219, 220, 221, 242, 245, 247, 253], "vmap": [1, 285, 287], "print": [1, 2, 3, 4, 6, 201, 202, 203, 209, 281, 284, 285, 286, 287], "ostream": 1, "os": [1, 6], "equival": [1, 27, 47, 58, 74, 96, 174, 218, 241, 243, 244], "check": [1, 6, 60, 231, 284], "bool": [1, 12, 14, 22, 23, 29, 30, 31, 32, 40, 41, 42, 43, 45, 56, 57, 59, 60, 106, 117, 119, 121, 122, 133, 135, 136, 141, 144, 145, 172, 184, 200, 211, 212, 213, 219, 220, 221, 222, 224, 228, 229, 231, 236, 239, 242, 244, 247, 251, 253], "is_equival": 1, "privat": 1, "fall": 1, "eval": [1, 2, 3, 4, 160, 209, 278, 280, 285, 287], "deriv": [1, 285], "base": [1, 106, 110, 112, 195, 197, 247, 253, 278, 280, 281, 284], "abov": [1, 3, 6, 134, 181, 194, 209, 284, 285, 289], "demonstr": [1, 286], "treat": [1, 60, 89, 90, 92, 93, 174], "paramet": [1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 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, 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, 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, 199, 200, 201, 202, 203, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 224, 225, 228, 229, 231, 236, 239, 240, 241, 242, 243, 244, 245, 247, 249, 251, 252, 253, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 274, 275, 278, 280, 285], "produc": [1, 242], "through": [1, 170, 196, 253, 286], "construct": [1, 4, 97, 128, 188], "its": [1, 6, 118, 132, 143, 160, 180, 190, 193, 194, 195, 203, 209, 244, 286, 289], "type": [1, 5, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 33, 50, 57, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 72, 73, 74, 75, 76, 77, 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, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 158, 159, 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, 197, 201, 209, 247, 253, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 284], "shape": [1, 3, 4, 47, 60, 61, 65, 66, 82, 85, 88, 91, 92, 93, 97, 103, 118, 128, 129, 136, 137, 138, 140, 141, 144, 145, 148, 175, 185, 187, 188, 189, 209, 211, 212, 213, 215, 216, 220, 222, 231, 268, 280, 284, 287, 289], "pass": [1, 3, 4, 47, 58, 130, 131, 183, 190, 201, 202, 209, 229, 239, 240, 241, 244, 249, 285], "re": [1, 4], "now": [1, 3, 6, 244, 286], "promot": 1, "dtype": [1, 3, 15, 26, 33, 57, 81, 97, 101, 106, 107, 128, 138, 140, 141, 144, 145, 180, 188, 205, 257, 284, 286, 287, 288], "promoted_dtyp": 1, "promote_typ": 1, "float32": [1, 15, 81, 101, 106, 107, 128, 138, 140, 144, 145, 180, 188, 205, 257, 284, 285, 286, 287, 288], "non": [1, 6, 223, 237, 269, 278], "point": [1, 2, 3, 6, 96, 135, 205], "out_dtyp": 1, "is_floating_point": 1, "cast": [1, 33, 91, 92, 93, 108, 224, 286], "up": [1, 3, 244], "determin": [1, 288], "x_cast": 1, "astyp": [1, 3, 224, 286], "y_cast": 1, "broadcasted_input": 1, "broadcast_arrai": 1, "out_shap": 1, "0": [1, 2, 3, 4, 6, 7, 15, 48, 53, 59, 64, 65, 66, 81, 94, 98, 106, 131, 136, 145, 147, 149, 165, 169, 180, 181, 182, 183, 184, 186, 191, 193, 194, 195, 196, 199, 200, 201, 209, 211, 212, 213, 214, 215, 216, 218, 219, 220, 221, 243, 246, 247, 248, 251, 252, 253, 255, 256, 257, 259, 260, 261, 267, 268, 270, 271, 272, 274, 278, 281, 284, 285, 286, 287, 288], "unique_ptr": 1, "make_uniqu": 1, "to_stream": 1, "handl": [1, 209], "resolv": 1, "No": [1, 3], "happen": [1, 3, 253, 280, 285], "alon": [1, 286], "effect": [1, 215, 285], "onli": [1, 3, 5, 6, 60, 65, 66, 106, 134, 205, 209, 228, 229, 231, 236, 239, 240, 241, 278, 288, 289], "execut": [1, 6, 286, 289], "depend": [1, 2, 57, 106, 284, 288, 289], "devic": [1, 5, 6, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 51, 53, 54, 55, 56, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 99, 100, 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, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 156, 158, 159, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 184, 187, 188, 189, 289, 290], "specifi": [1, 15, 33, 66, 89, 90, 97, 98, 106, 107, 124, 128, 137, 147, 173, 174, 175, 178, 179, 183, 186, 188, 211, 252, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 274, 289], "memori": [1, 5, 160, 278, 285, 286], "ha": [1, 3, 4, 5, 57, 88, 89, 91, 92, 93, 98, 137, 211, 222, 278, 280, 284, 285, 287, 289], "been": [1, 3, 285], "try": [1, 6], "naiv": 1, "gener": [1, 2, 15, 81, 89, 90, 107, 136, 140, 141, 144, 145, 253, 281, 284, 285, 290], "version": [1, 6, 72, 113, 117, 134, 163, 186, 281, 284], "declar": 1, "member": [1, 209, 234, 237], "method": [1, 3, 7, 8, 9, 26, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 209, 278], "each": [1, 50, 72, 78, 118, 131, 134, 135, 137, 147, 154, 155, 165, 179, 186, 187, 215, 216, 217, 219, 247, 253, 259, 281, 285], "find": [1, 2, 6], "pointwis": 1, "captur": [1, 209], "templat": 1, "axpby_impl": 1, "typenam": 1, "t": [1, 3, 76, 135, 183, 191, 192, 193, 194, 195, 196, 199, 200, 209, 289], "readi": 1, "fill": [1, 97, 129, 180, 189], "malloc_or_wait": 1, "synchron": 1, "avail": [1, 2, 3, 4, 6, 8, 205, 289], "There": [1, 209], "wait": [1, 3], "here": [1, 3, 285, 288, 289], "request": 1, "pressur": 1, "condit": [1, 187, 289], "set_data": 1, "nbyte": 1, "collect": [1, 198, 202, 283], "pointer": 1, "x_ptr": 1, "y_ptr": 1, "out_ptr": 1, "relev": 1, "static_cast": 1, "size_t": 1, "out_idx": 1, "size": [1, 3, 4, 50, 66, 72, 80, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 97, 101, 106, 134, 135, 137, 148, 165, 168, 209, 212, 213, 217, 220, 244, 285, 286], "map": [1, 4, 108, 202, 217, 224], "linear": [1, 3, 4, 5, 202, 209, 218, 231, 244, 246, 248, 250, 254, 255, 256, 271, 272, 273, 278], "indic": [1, 13, 22, 23, 24, 25, 98, 165, 174, 175, 183, 236, 238, 259, 284], "offset": [1, 3], "x_offset": 1, "elem_to_loc": 1, "stride": [1, 65, 66, 212, 213, 247, 284], "y_offset": 1, "contigu": 1, "regularli": 1, "default": [1, 6, 12, 14, 15, 22, 23, 24, 25, 60, 64, 65, 66, 72, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 98, 101, 106, 107, 108, 117, 119, 121, 122, 128, 132, 133, 134, 135, 136, 137, 138, 140, 141, 143, 144, 145, 147, 148, 149, 164, 165, 168, 169, 172, 178, 179, 180, 181, 182, 183, 184, 186, 188, 191, 192, 193, 194, 195, 196, 198, 199, 200, 205, 211, 212, 213, 220, 222, 224, 229, 231, 236, 239, 242, 243, 244, 247, 251, 253, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 278, 281, 283, 286, 288, 290], "row": [1, 81, 101, 134, 180], "major": 1, "henc": [1, 134], "doesn": [1, 209], "addit": [1, 3, 11, 211, 219, 221, 242, 245, 278], "abl": [1, 134], "all": [1, 4, 6, 13, 24, 66, 78, 81, 84, 87, 90, 93, 118, 131, 132, 168, 197, 209, 224, 225, 229, 232, 233, 234, 237, 239, 242, 244, 251, 253, 278, 281, 284, 285, 287, 290], "incom": 1, "accordingli": 1, "dispatch": 1, "float16": [1, 108, 205, 224, 285, 286], "bfloat16": [1, 286], "complex64": 1, "throw": 1, "error": [1, 6, 76, 77, 165, 218, 244, 254, 255, 256, 264, 265, 286], "encount": 1, "unexpect": [1, 15], "regist": [1, 4], "op": [1, 130, 229, 285], "assert": 1, "2": [1, 2, 3, 4, 66, 76, 83, 86, 88, 89, 90, 91, 92, 93, 106, 112, 118, 134, 143, 178, 180, 181, 182, 191, 192, 193, 194, 199, 205, 209, 213, 218, 245, 251, 255, 261, 267, 268, 278, 284, 285, 286, 287, 288, 289], "1": [1, 3, 4, 15, 24, 25, 65, 66, 82, 83, 85, 86, 88, 89, 90, 91, 92, 93, 94, 102, 106, 118, 130, 132, 134, 137, 145, 158, 164, 174, 183, 191, 192, 193, 194, 195, 196, 199, 200, 205, 209, 211, 212, 213, 214, 215, 216, 218, 219, 220, 221, 222, 243, 245, 247, 248, 251, 252, 255, 256, 257, 258, 259, 260, 261, 262, 264, 266, 267, 268, 272, 274, 278, 280, 284, 286, 287, 288, 289], "correct": [1, 193, 194, 195, 284, 285], "els": [1, 3, 209, 229, 285], "float16_t": 1, "bfloat16_t": 1, "complex64_t": 1, "runtime_error": 1, "support": [1, 3, 5, 6, 13, 65, 66, 108, 118, 134, 284, 286, 288], "have": [1, 3, 6, 60, 89, 90, 92, 93, 118, 137, 196, 201, 242, 249, 283, 284, 285, 289], "rememb": 1, "3": [1, 3, 6, 106, 196, 281, 284, 286, 287], "complic": 1, "keep": [1, 12, 14, 22, 23, 117, 119, 121, 122, 133, 172, 184, 209, 228, 285], "mind": [1, 3], "half": [1, 15, 141, 145, 247, 285], "precis": [1, 3, 209, 218], "direct": [1, 3, 196, 226, 289], "fix": [1, 3, 6, 285], "possibl": [1, 3, 118, 165, 217, 284, 289], "due": 1, "transpos": [1, 3, 27, 135], "aren": 1, "guarante": 1, "fit": [1, 134, 289], "requir": [1, 3, 209, 285, 286], "column": [1, 81, 101, 134], "inplac": 1, "expect": [1, 3, 212, 213, 214, 215, 216, 251, 253, 284], "answer": 1, "copi": [1, 3, 5, 132, 164, 286], "simpli": [1, 3, 6, 246, 271, 278], "catlas_saxpbi": 1, "axpby_impl_acceler": 1, "first": [1, 2, 3, 4, 6, 94, 98, 114, 116, 118, 132, 143, 173, 178, 183, 193, 194, 195, 201, 209, 219, 258, 286, 289], "mode": [1, 67, 227, 236, 238], "i": [1, 3, 103, 106, 194, 209, 212, 213, 215, 216, 229, 264], "e": [1, 4, 6, 76, 103, 158, 192, 211, 212, 213, 215, 216, 219, 220, 221, 229, 245, 275, 280, 285, 290], "match": [1, 6, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 231, 284, 286], "transposit": 1, "data_s": 1, "items": 1, "flag": [1, 286], "copy_inplac": 1, "copytyp": 1, "n": [1, 3, 26, 65, 66, 81, 82, 84, 85, 87, 88, 91, 93, 101, 180, 184, 211, 212, 213, 215, 216, 264, 268], "incx": 1, "inci": 1, "great": 1, "But": [1, 289], "criteria": 1, "luckili": [1, 285], "alwai": [1, 201], "With": 1, "final": [1, 2, 3, 4], "singl": [1, 4, 78, 103, 108, 131, 185, 284, 288], "row_contigu": 1, "col_contigu": 1, "common": [1, 285], "hit": 1, "mileston": 1, "enough": [1, 285], "run": [1, 3, 4, 5, 6, 130, 160, 191, 193, 194, 195, 211, 224, 285, 289, 290], "If": [1, 3, 6, 12, 14, 15, 22, 23, 24, 25, 57, 60, 63, 64, 67, 78, 91, 92, 93, 96, 97, 98, 106, 108, 117, 118, 119, 121, 122, 128, 131, 132, 133, 137, 147, 163, 164, 165, 172, 174, 175, 178, 183, 184, 186, 188, 202, 211, 212, 213, 219, 221, 222, 229, 231, 239, 244, 247, 249, 251, 268, 285, 288, 289, 290], "plan": 1, "stop": [1, 3, 15, 107, 170, 284], "enjoi": 1, "speed": 1, "appl": [1, 3, 5, 6, 289], "silicon": [1, 3, 5, 6, 289], "address": 1, "shade": 1, "languag": [1, 205], "kernel": [1, 65, 66, 284], "written": 1, "help": [1, 3, 289], "resourc": 1, "walkthrough": 1, "pipelin": 1, "specif": [1, 6], "cpp": 1, "algorithm": [1, 196], "launch": [1, 284], "exactli": [1, 3, 231], "mani": [1, 165, 212, 213, 217, 285], "thread": 1, "pick": 1, "updat": [1, 2, 3, 4, 194, 196, 200, 202, 211, 224, 231, 241, 280, 285], "assign": [1, 278], "axpby_gener": 1, "buffer": [1, 286], "constant": [1, 3, 6, 131, 199, 211, 219, 221, 245, 268, 286], "4": [1, 3, 72, 106, 134, 135, 154, 205, 211, 220, 244, 253, 284, 287, 289], "5": [1, 2, 3, 6, 106, 136, 199, 211, 214, 215, 216, 220, 267, 284], "x_stride": 1, "6": [1, 3, 106, 154, 199, 253, 255, 256, 268, 284, 287], "y_stride": 1, "7": [1, 3, 106, 134, 284], "ndim": [1, 106], "8": [1, 3, 6, 106, 134, 191, 192, 193, 194, 195, 199, 205, 220, 253, 258, 284, 287, 289], "uint": 1, "index": [1, 5, 7, 9, 24, 80, 81, 98, 132, 174, 175, 183], "thread_position_in_grid": 1, "convert": [1, 57, 244, 285, 286, 287], "instanti": [1, 4, 285], "uniqu": [1, 281], "host": 1, "name": [1, 108, 134, 135, 152, 153, 154, 155, 198, 209, 219, 228, 231, 233, 284, 288], "identifi": [1, 201, 283], "instantiate_axpbi": 1, "type_nam": 1, "host_nam": 1, "axpby_general_": 1, "compil": [1, 6, 285], "mlx_ext": 1, "metallib": [1, 6], "see": [1, 3, 4, 6, 8, 28, 29, 30, 31, 32, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 51, 53, 54, 55, 56, 58, 59, 106, 152, 153, 209, 211, 215, 218, 227, 243, 244, 247, 248, 251, 254, 255, 256, 272, 284, 287, 289], "later": [1, 6], "co": [1, 251], "locat": [1, 240, 241, 289], "share": [1, 5, 72, 134, 135], "register_librari": 1, "potenti": 1, "path": [1, 6, 154, 155, 231], "tri": 1, "load": [1, 4, 5, 231], "hasn": 1, "alreadi": [1, 3], "static": [1, 6], "object": [1, 8, 26, 37, 57, 136, 141, 144, 145, 186, 201, 202, 215, 283], "why": [1, 3], "packag": [1, 2, 4], "process": [1, 3, 67, 202, 216, 217, 253, 283], "logic": [1, 114, 115, 116], "grid": 1, "shown": 1, "below": [1, 6, 106, 180, 182, 205, 285], "prepar": [1, 3], "carri": 1, "should": [1, 2, 3, 4, 6, 103, 134, 160, 175, 183, 185, 201, 209, 212, 213, 215, 216, 236, 242, 249, 260, 278, 283, 285, 286, 290], "d": [1, 3, 102, 106, 118, 130, 174, 180, 181, 182, 191, 193, 195, 203, 216, 289], "ostringstream": 1, "kname": 1, "axpby_": 1, "general_": 1, "type_to_nam": 1, "make": [1, 3, 4, 6, 118, 209, 285, 287, 289], "sure": [1, 3, 6, 209], "look": [1, 3], "folder": 1, "get_colocated_mtllib_path": 1, "get_kernel": 1, "str": [1, 67, 98, 106, 108, 151, 152, 153, 154, 155, 183, 201, 203, 224, 225, 228, 229, 231, 233, 235, 239, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268], "encod": [1, 247, 251, 253], "compute_encod": 1, "get_command_encod": 1, "setcomputepipelinest": 1, "those": [1, 3, 209], "nelem": 1, "set_array_buff": 1, "setbyt": 1, "sizeof": 1, "threadgroup": 1, "higher": [1, 102], "than": [1, 3, 57, 67, 74, 99, 100, 104, 105, 118, 196, 202, 247, 252, 267, 274, 289], "max": [1, 106, 120, 195, 243, 258, 260, 268, 270, 289], "allow": [1, 197, 209, 241, 278, 284, 287], "tgp_size": 1, "min": [1, 106, 123, 243, 270], "maxtotalthreadsperthreadgroup": 1, "3d": [1, 211, 216], "mtl": 1, "group_dim": 1, "grid_dim": 1, "divid": [1, 96, 134], "among": 1, "dispatchthread": 1, "few": [1, 3, 4, 5, 160, 285, 287], "thing": [1, 3], "note": [1, 3, 6, 13, 65, 66, 89, 90, 106, 134, 137, 209, 286, 288], "befor": [1, 3, 6, 24, 132, 228, 253, 284, 285], "move": [1, 124, 289], "track": [1, 209, 211], "activ": [1, 6, 215, 223, 252, 253, 269, 274, 275], "command": [1, 6], "instead": [1, 6, 209, 241, 251, 285], "end_encod": 1, "end": [1, 134, 248, 252, 261, 267, 272, 274], "until": [1, 285, 287], "limit": [1, 63, 284], "flush": 1, "enqueu": 1, "commit": 1, "associ": [1, 154, 155, 285], "suggest": 1, "deeper": 1, "dive": 1, "studi": 1, "come": [1, 3], "far": [1, 280], "built": [1, 6, 285], "includ": [1, 225, 244, 284, 287, 288, 290], "forward": [1, 183, 285], "diff": 1, "push": 1, "along": [1, 22, 23, 64, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 106, 147, 163, 165, 169, 174, 175, 178, 209], "similarli": [1, 6, 118, 285], "scale_arr": 1, "contribut": 1, "tangent_x": 1, "tangent_i": 1, "revers": [1, 179, 251], "arg": [1, 3, 8, 47, 58, 78, 154, 155, 160], "push_back": 1, "fulli": [1, 5, 286, 289], "overal": 1, "directori": [1, 3, 6], "extens": [1, 108, 205, 288], "h": [1, 65, 66, 106, 211, 213, 215, 216, 285], "mlx_sample_extens": 1, "__init__": [1, 3, 4, 7, 8, 9, 26, 209, 278], "py": [1, 3, 6], "cmakelist": 1, "txt": 1, "setup": [1, 2, 4, 6], "hold": [1, 3, 8, 106, 197], "instal": 1, "pybind11": [1, 6], "sinc": [1, 3, 4, 196, 278, 286, 289], "compon": [1, 3], "etc": [1, 134, 209], "becom": 1, "pybind11_modul": 1, "m": [1, 6, 81, 106, 180, 191], "doc": [1, 4], "sampl": [1, 2, 3, 107, 136, 137, 138, 141, 144, 145, 268, 281], "_a": 1, "pos_onli": 1, "kw_onli": 1, "none": [1, 3, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 51, 53, 54, 55, 56, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 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, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 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, 186, 187, 188, 189, 201, 202, 218, 224, 228, 229, 239, 242, 251, 253, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 284], "r": [1, 3, 183, 215], "pbdoc": 1, "most": [1, 137, 209, 284, 285], "complex": [1, 89, 90, 91, 92, 93, 136, 141, 144, 145, 201, 209, 241], "bell": 1, "whistl": 1, "liter": 1, "string": [1, 286, 288], "modul": [1, 3, 4, 190, 244, 249, 253, 283, 285], "ensur": [1, 6, 264], "caster": 1, "find_packag": 1, "config": 1, "add_librari": 1, "sourc": [1, 124, 179], "target_sourc": 1, "cmake_current_list_dir": 1, "header": 1, "target_include_directori": 1, "target_link_librari": 1, "attach": 1, "conveni": [1, 4], "mlx_build_metallib": 1, "target": [1, 183, 257, 259, 260, 261, 262, 263, 264, 265, 266, 267], "destin": [1, 124], "automat": [1, 5, 108, 287, 288, 289], "practic": 1, "mlx_build_met": [1, 6], "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": 1, "build_shared_lib": 1, "target_link_opt": 1, "wl": 1, "rpath": 1, "loader_path": 1, "onc": [1, 160], "describ": [1, 285], "util": [1, 3, 5, 6, 154, 209], "__name__": [1, 3], "__main__": [1, 3], "descript": [1, 3, 205], "ext_modul": 1, "cmakeextens": 1, "cmdclass": 1, "build_ext": 1, "cmakebuild": 1, "package_dir": 1, "package_data": 1, "dylib": 1, "zip_saf": 1, "fals": [1, 3, 12, 14, 22, 23, 29, 30, 31, 32, 40, 41, 42, 43, 45, 56, 59, 60, 106, 117, 119, 121, 122, 133, 172, 184, 187, 200, 201, 202, 205, 219, 220, 222, 229, 231, 239, 242, 244, 247, 251, 253, 286], "python_requir": 1, "even": [1, 3, 285, 286], "though": [1, 3, 285, 286], "j8": 1, "libmlx_ext": 1, "cpython": 1, "3x": 1, "darwin": 1, "pip": [1, 6], "after": [1, 3, 4, 24, 96, 132, 134, 211, 219, 221, 242, 253, 267, 289], "plai": [1, 3], "ones": [1, 3, 129, 154, 160, 180, 240, 241, 244, 284], "b": [1, 3, 11, 13, 60, 73, 74, 75, 96, 99, 100, 102, 104, 105, 106, 113, 114, 116, 118, 120, 123, 125, 130, 134, 171, 178, 183, 222, 284, 285, 286, 287, 288, 289], "f": [1, 2, 4, 106, 194, 209, 286], "item": [1, 2, 3, 4, 202, 285, 286, 287], "true": [1, 2, 3, 60, 106, 135, 163, 187, 201, 202, 205, 209, 211, 212, 213, 219, 220, 221, 222, 228, 229, 231, 236, 239, 244, 247, 251, 253, 264], "quick": [1, 5], "benchmark": 1, "compar": [1, 60], "time": [1, 3, 6, 160, 209, 285, 289], "set_default_devic": 1, "256": [1, 4], "512": [1, 3, 253, 289], "random": [1, 2, 3, 4, 5, 211, 220, 231, 236, 289, 290], "normal": [1, 2, 3, 144, 198, 211, 219, 220, 221, 245, 253, 286, 289], "bench": 1, "warm": 1, "rang": [1, 2, 3, 4, 6, 15, 107, 255, 256, 280, 281, 285, 289], "100": [1, 2, 3, 285, 289], "5000": 1, "simple_tim": 1, "custom_tim": 1, "3f": [1, 4], "custom": [1, 253], "114": 1, "109": 1, "modest": 1, "improv": [1, 3, 191, 192, 193, 194, 195, 199], "awai": [1, 3], "good": [1, 6, 289], "nn": [1, 3, 4, 154, 202, 209, 278, 280, 285], "grad": [1, 2, 4, 183, 280, 284, 285, 287], "simplifi": [1, 285], "full": [1, 4, 47, 58, 67, 163, 240, 241, 285], "implement": [2, 4, 106, 191, 192, 193, 194, 195, 196, 197, 198, 199, 217, 228, 242, 247, 249, 251, 252, 253, 274, 286], "basic": [2, 149], "model": [2, 4, 5, 154, 190, 202, 209, 224, 227, 229, 231, 235, 236, 238, 239, 240, 242, 253, 278, 280, 285], "problem": [2, 4, 209], "metadata": 2, "num_featur": [2, 211], "num_exampl": 2, "1_000": 2, "num_it": 2, "10_000": 2, "iter": [2, 4, 202, 281, 285], "sgd": [2, 4, 196, 280], "lr": [2, 196], "01": [2, 194], "rate": [2, 191, 192, 193, 194, 195, 196, 199, 200], "ll": [2, 4, 261], "synthet": 2, "dataset": [2, 285], "matrix": [2, 72, 81, 101, 106, 118, 134, 135, 244], "ground": [2, 3, 267], "truth": [2, 267], "w_star": 2, "valu": [2, 3, 10, 15, 22, 23, 37, 57, 60, 63, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 97, 106, 107, 131, 136, 137, 138, 140, 141, 144, 145, 174, 175, 183, 186, 190, 194, 198, 201, 202, 205, 214, 215, 216, 220, 222, 228, 242, 243, 247, 252, 253, 257, 258, 259, 260, 261, 263, 264, 265, 266, 267, 274, 278], "gaussian": [2, 218, 254, 255, 256], "nois": 2, "exampl": [2, 3, 4, 15, 106, 174, 209, 211, 220, 229, 231, 236, 239, 257, 280, 281, 284, 285, 286, 287, 288], "noisi": 2, "label": [2, 259], "ep": [2, 191, 192, 193, 194, 195, 199, 211, 219, 220, 221, 245, 258, 268], "1e": [2, 4, 13, 191, 192, 193, 194, 195, 199, 211, 219, 220, 221, 245, 258, 268], "us": [2, 3, 4, 5, 6, 15, 72, 74, 94, 106, 118, 134, 135, 147, 148, 191, 193, 194, 195, 196, 197, 201, 209, 215, 217, 218, 222, 224, 228, 240, 241, 242, 244, 247, 251, 253, 255, 256, 258, 278, 280, 281, 283, 284, 287, 289], "weight": [2, 65, 66, 194, 196, 200, 202, 209, 231, 235, 244, 259, 278, 285], "squar": [2, 3, 101, 150, 166, 183, 191, 193, 194, 195, 202, 209, 245, 265, 267, 286], "loss": [2, 4, 183, 209, 280, 285], "loss_fn": [2, 4, 280], "w": [2, 66, 72, 134, 135, 183, 200, 211, 213, 215, 216, 222], "mean": [2, 3, 4, 183, 209, 211, 219, 229, 245, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 286], "grad_fn": 2, "initi": [2, 3, 209, 211, 219, 220, 221, 222, 243, 245, 278, 285], "randomli": [2, 3, 214, 215, 216], "Then": [2, 6], "repeatedli": 2, "_": [2, 3, 209, 281, 285, 289], "verifi": [2, 6], "close": [2, 5, 6, 13], "error_norm": 2, "5f": 2, "someth": [2, 3, 284], "00005": 2, "00364": 2, "complet": [2, 3, 6, 240, 241, 289], "logist": [2, 158, 250, 255, 256, 273], "github": [2, 4, 6], "repo": [2, 4, 6], "enabl": [3, 6, 200], "larg": [3, 209, 242, 264, 285], "ish": 3, "transform": [3, 5, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 190, 209, 211, 219, 221, 222, 228, 229, 239, 244, 247, 284], "compromis": 3, "eas": 3, "llama": 3, "famili": 3, "less": [3, 24, 105, 132, 247, 267], "200": 3, "line": [3, 285, 286], "python": [3, 37, 50, 57, 78, 201, 202, 203, 278, 283, 286], "neural": [3, 5, 199, 217, 223, 269, 278], "network": [3, 5, 199, 211, 215, 217, 278], "build": [3, 5, 278], "concis": 3, "architectur": [3, 209, 241, 289], "notabl": [3, 5], "rope": [3, 209], "posit": [3, 24, 98, 124, 132, 183, 202, 209, 212, 213, 242, 247, 251, 268], "option": [3, 12, 14, 15, 22, 23, 24, 25, 26, 31, 32, 64, 65, 66, 67, 72, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 97, 98, 101, 106, 107, 108, 117, 119, 121, 122, 128, 131, 132, 133, 134, 135, 136, 137, 138, 140, 141, 143, 144, 145, 147, 148, 163, 164, 165, 168, 169, 172, 174, 175, 178, 179, 180, 181, 182, 183, 184, 186, 188, 191, 192, 193, 194, 195, 196, 199, 200, 201, 202, 211, 212, 213, 222, 224, 228, 229, 231, 239, 242, 244, 247, 251, 253, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 281, 288, 290], "kei": [3, 136, 137, 138, 140, 141, 143, 144, 145, 198, 201, 202, 228, 229, 239, 242, 247, 281, 283], "cach": [3, 247], "concaten": 3, "project": [3, 242], "llamaattent": 3, "self": [3, 4, 7, 9, 26, 27, 28, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 51, 53, 54, 55, 56, 57, 58, 59, 209, 223, 269, 278], "dim": [3, 178, 217, 219, 220, 221, 242, 245, 247, 251, 253], "num_head": [3, 242, 253], "super": [3, 4, 209, 278], "tradit": [3, 215, 216, 247], "query_proj": 3, "bia": [3, 72, 134, 135, 193, 194, 195, 202, 209, 212, 213, 222, 229, 231, 239, 242, 244], "key_proj": 3, "value_proj": 3, "out_proj": [3, 278], "__call__": [3, 4, 209, 278], "queri": [3, 242], "mask": [3, 236, 242, 284], "extract": [3, 209, 228, 278], "l": [3, 4, 209, 211, 212, 267], "reshap": [3, 106, 284], "combin": 3, "key_cach": 3, "value_cach": 3, "sqrt": [3, 76, 191, 192, 193, 194, 199, 211, 219, 220, 221, 222, 245, 251], "score": 3, "softmax": [3, 259], "values_hat": 3, "rm": 3, "swiglu": 3, "rmsnorm": [3, 209], "llamaencoderlay": 3, "mlp_dim": [3, 253], "norm1": 3, "norm2": 3, "linear1": 3, "linear2": 3, "linear3": 3, "sigmoid": [3, 250, 255, 256, 257, 273], "instanc": [3, 134, 203, 209, 220, 224, 225, 226, 229, 232, 233, 239, 241, 249, 278, 286], "embed": [3, 209, 247, 251, 258], "emb": [3, 217, 251], "token": [3, 217], "num_lay": [3, 4, 280], "vocab_s": 3, "norm": [3, 195, 196, 219, 268], "multiheadattent": [3, 209], "create_additive_causal_mask": 3, "list": [3, 8, 12, 14, 26, 29, 30, 40, 41, 42, 43, 45, 50, 53, 56, 57, 59, 61, 64, 78, 80, 83, 84, 86, 87, 89, 90, 92, 93, 97, 98, 103, 106, 117, 119, 121, 122, 128, 131, 133, 136, 137, 138, 140, 141, 144, 145, 148, 163, 165, 168, 169, 172, 178, 179, 183, 184, 185, 188, 193, 194, 195, 196, 201, 203, 209, 229, 231, 232, 233, 234, 237, 239, 240, 241, 278, 283, 285], "still": [3, 6, 106, 285], "consid": [3, 13, 60, 201, 202, 219, 283], "train": [3, 4, 209, 211, 214, 215, 216, 227, 229, 239], "ignor": [3, 63], "whatsoev": 3, "rest": [3, 202, 247], "subsect": 3, "prompt": 3, "autoregress": 3, "yield": [3, 4, 281], "temp": 3, "causal": 3, "save": [3, 5, 108, 134, 152, 153, 154, 155, 235, 285], "append": [3, 118, 285], "store": 3, "per": [3, 4, 72, 134, 135, 197, 211, 219, 220, 221, 245, 285], "care": [3, 285], "last": [3, 25, 57, 84, 87, 89, 90, 92, 93, 94, 102, 118, 137, 164, 178, 212, 213, 215, 216, 219, 286], "logit": [3, 137, 257, 259], "next": [3, 4], "categor": 3, "lazili": [3, 209], "noth": [3, 209, 285], "yet": [3, 106, 209, 278, 284, 285, 287], "forc": [3, 4, 209, 287], "choos": [3, 247], "pars": 3, "feed": 3, "loop": [3, 4, 285], "unsqueez": 3, "sequenc": [3, 211, 212, 253, 281, 289], "length": [3, 168, 211, 212], "len": [3, 84, 87, 90, 93], "overwrit": 3, "discard": [3, 201], "old": 3, "moment": [3, 193, 194, 195], "anymor": 3, "everyth": 3, "small": [3, 160, 211, 219, 221, 245, 264, 268, 289], "10": [3, 4, 110, 149, 154, 160, 202, 209, 231, 284], "12": 3, "8192": 3, "1024": 3, "actual": [3, 15, 231, 278, 285], "materi": [3, 5], "could": [3, 209], "20_000": 3, "machin": [3, 5, 6, 199], "8gb": 3, "ram": 3, "32": [3, 4, 134, 135, 205], "44": 3, "doubl": 3, "bracket": 3, "becaus": [3, 209, 285], "batch": [3, 118, 211, 212, 213, 215, 216, 242, 285], "zip": [3, 4], "haven": 3, "anyth": [3, 183, 285], "result": [3, 15, 57, 72, 102, 106, 108, 118, 130, 135, 147, 149, 169, 178, 187, 202, 251, 286], "similar": [3, 202, 240, 241, 242, 258, 286, 288], "runtim": 3, "section": [3, 6, 165, 268], "access": [3, 37, 209, 278, 285, 289], "origin": [3, 191, 192, 193, 195, 196, 211, 286], "sentencepiec": 3, "pytorch": [3, 5, 219], "compat": [3, 137, 288], "npz": [3, 108, 154, 155, 231, 235, 288], "file": [3, 6, 108, 151, 152, 153, 154, 155, 231, 235, 288], "directli": 3, "argpars": 3, "itertool": [3, 202], "starmap": [3, 202], "np": [3, 4, 286, 287], "torch": [3, 286], "map_torch_to_mlx": 3, "tok_embed": 3, "elif": 3, "replac": [3, 240, 241, 253, 267], "attention_norm": 3, "ffn_norm": 3, "wq": 3, "wk": 3, "wv": 3, "wo": 3, "w1": 3, "w2": 3, "w3": 3, "ffn": 3, "separ": [3, 47, 58, 219], "submodul": [3, 4, 209, 229, 230, 239, 241], "feed_forward": 3, "parser": 3, "argumentpars": 3, "add_argu": 3, "torch_weight": 3, "output_fil": 3, "parse_arg": 3, "state": [3, 4, 197, 198, 209, 280, 281], "savez": [3, 288], "k": [3, 81, 180, 181, 182, 222, 229], "v": [3, 67, 209, 229, 286], "left": [3, 106, 134, 218, 247, 255, 256, 261, 268], "disk": 3, "text": [3, 196, 223, 248, 252, 260, 261, 264, 267, 268, 269, 270, 272, 274], "format": [3, 108, 151, 152, 153, 154, 155, 286], "oper": [3, 5, 33, 160, 163, 170, 175, 196, 209, 253, 284, 285, 286, 287, 289, 290], "dictionari": [3, 152, 153, 197, 198, 201, 209, 228, 240, 241, 283, 288], "represent": [3, 134, 201, 203], "tree_unflatten": 3, "helper": 3, "weight_fil": 3, "incur": 3, "sever": [3, 65, 66, 154, 155, 288], "futur": [3, 244, 284, 285], "pth": 3, "current": [3, 5, 6, 65, 66, 134, 209, 285], "around": 3, "m1": [3, 289], "ultra": 3, "7b": 3, "me": 3, "ishmael": 3, "year": 3, "ago": 3, "never": [3, 285], "long": 3, "info": [3, 6], "247": 3, "press": [3, 106], "enter": 3, "littl": 3, "monei": 3, "my": [3, 6], "purs": 3, "greater": [3, 24, 100, 132, 252, 274], "consequ": 3, "walk": 3, "down": 3, "gower": 3, "street": 3, "afternoon": 3, "heavi": 3, "rain": 3, "saw": 3, "off": [3, 6, 285], "man": 3, "rag": 3, "who": 3, "sat": 3, "upon": [3, 202], "hi": 3, "bundl": 3, "hard": 3, "wet": 3, "he": 3, "were": [3, 289], "cry": 3, "watch": 3, "him": 3, "observ": 3, "numer": [3, 106, 113, 117, 163, 191, 192, 193, 194, 195, 199, 211, 219, 220, 221, 245, 258, 268, 285], "crowd": 3, "wa": [3, 198, 285], "hurri": 3, "437": 3, "330": 3, "second": [3, 114, 116, 118, 173, 183, 193, 194, 195, 258, 289], "spent": 3, "amount": 3, "39": 3, "ms": 3, "By": [3, 286], "bigger": 3, "remain": [3, 183, 214, 215, 216], "almost": 3, "nobodi": 3, "took": 3, "least": [3, 63, 134], "notic": [3, 288], "distanc": [3, 268], "had": 3, "doubt": 3, "minut": 3, "straight": 3, "slowli": 3, "rais": [3, 106, 165, 231], "ey": 3, "speak": [3, 106], "resum": 3, "postur": 3, "stood": 3, "feel": 3, "pain": 3, "heart": 3, "smile": 3, "face": 3, "am": 3, "someon": 3, "three": 3, "quarter": 3, "hour": 3, "made": 3, "immedi": [3, 224], "repli": 3, "again": [3, 209], "hand": [3, 285], "did": 3, "accustom": 3, "thu": [3, 209], "question": [3, 285], "reason": [3, 284], "tell": [3, 286], "understand": 3, "579": 3, "690": 3, "num": [3, 107, 143], "500": [3, 289], "628": 3, "went": 3, "nervou": 3, "trembl": 3, "told": 3, "And": 3, "perhap": 3, "surpris": 3, "matter": [3, 209], "shall": 3, "anyhow": 3, "friend": 3, "ye": 3, "slight": [3, 285], "kind": 3, "longer": [3, 67], "soon": 3, "unless": [3, 106, 278], "unlik": [3, 13, 215, 216], "strang": 3, "amus": 3, "That": 3, "secret": 3, "disappoint": 3, "mine": 3, "cannot": [3, 63, 284, 286], "happi": 3, "ask": 3, "Is": [3, 251, 253], "shop": 3, "bui": 3, "food": 3, "633": 3, "21": 3, "475": 3, "su": 3, "j": [3, 6, 106, 192, 193, 195, 215], "lu": 3, "pan": 3, "murtadha": 3, "wen": 3, "liu": 3, "2021": 3, "roform": [3, 247], "enhanc": [3, 247, 285], "rotari": [3, 247], "arxiv": [3, 191, 196, 219, 220, 221, 223, 245, 269], "preprint": [3, 191, 196], "2104": 3, "09864": 3, "zhang": 3, "sennrich": 3, "2019": [3, 194], "root": [3, 150, 166, 245], "advanc": 3, "inform": [3, 4, 152, 153, 209, 211, 218, 242, 289], "system": 3, "shazeer": 3, "2020": 3, "glu": 3, "variant": [3, 195, 267], "2002": 3, "05202": 3, "classifi": 4, "mnist": 4, "As": [4, 174, 209], "mlp": [4, 209, 253, 280], "inherit": [4, 283], "standard": [4, 37, 57, 118, 138, 253, 287], "idiom": 4, "input_dim": [4, 209, 222, 244], "hidden_dim": [4, 278, 280], "output_dim": [4, 209, 222, 244], "layer_s": 4, "idim": 4, "odim": 4, "maximum": [4, 22, 63, 209, 246, 251, 255, 256, 271, 278, 285], "cross": [4, 257, 259], "entropi": [4, 257, 259], "sub": [4, 143], "commonli": [4, 240], "cross_entropi": [4, 209], "accuraci": 4, "valid": [4, 67, 186, 201, 229, 239, 283], "eval_fn": 4, "argmax": 4, "loader": 4, "num_class": [4, 280], "batch_siz": [4, 280], "num_epoch": [4, 280], "learning_r": [4, 191, 192, 193, 194, 195, 196, 199, 200, 280], "train_imag": [4, 280], "train_label": [4, 280], "test_imag": 4, "test_label": 4, "shuffl": 4, "minibatch": 4, "batch_iter": [4, 280], "perm": 4, "permut": 4, "id": [4, 6], "put": 4, "trainabl": [4, 190, 209, 278], "loss_and_grad_fn": [4, 280], "value_and_grad": [4, 209, 240, 278, 280, 286, 287], "epoch": 4, "test": [4, 6], "confus": 4, "decent": 4, "95": 4, "brought": 5, "research": 5, "except": [5, 81, 88, 89, 91, 92, 93, 219, 231, 284, 286], "featur": [5, 65, 66, 211, 219, 220, 221, 222, 244, 245, 247, 253, 285], "main": [5, 81, 202, 209], "differ": [5, 171, 267], "lazi": [5, 278, 287], "multi": [5, 212, 213, 284, 286], "cpu": [5, 289], "gpu": [5, 284, 289], "inspir": 5, "jax": [5, 281], "arrayfir": 5, "unifi": 5, "live": [5, 289], "guid": 5, "convers": 5, "regress": [5, 264], "layer": [5, 209, 215, 216, 219, 221, 222, 236, 241, 244, 249, 253, 275, 278], "perceptron": 5, "llm": 5, "infer": [5, 97, 108], "fft": 5, "algebra": 5, "tree": [5, 78, 98, 160, 183, 186, 197, 201, 202, 203], "develop": [5, 6], "document": [5, 47, 58, 152, 153, 284], "meet": 6, "seri": 6, "chip": 6, "nativ": 6, "maco": 6, "13": 6, "recommend": [6, 196], "14": 6, "sonoma": 6, "distribut": [6, 136, 137, 138, 140, 144, 145, 222, 262, 266, 268], "probabl": [6, 141, 214, 215, 216, 244, 262, 289], "platform": 6, "processor": 6, "arm": [6, 205], "i386": 6, "switch": 6, "conda": 6, "17": 6, "g": [6, 106, 134, 199, 200, 275, 285, 290], "clang": 6, "cmake": 6, "24": 6, "xcode": 6, "15": [6, 106], "environ": 6, "via": [6, 285, 286], "rosetta": 6, "unam": 6, "p": [6, 136, 193, 195, 209, 214, 215, 216, 268], "clone": 6, "git": 6, "com": 6, "ml": 6, "explor": 6, "cd": 6, "brew": 6, "global": [6, 142, 281], "env": 6, "cmake_build_parallel_level": 6, "edit": [6, 241], "unittest": 6, "discov": 6, "stub": 6, "dev": 6, "generate_stub": 6, "mkdir": 6, "either": [6, 11, 47, 57, 58, 63, 73, 74, 75, 96, 99, 100, 104, 105, 106, 113, 118, 120, 123, 125, 171, 183, 249], "libmlx": 6, "preprocessor": 6, "metal_path": 6, "mlx_build_test": 6, "ON": 6, "mlx_build_exampl": 6, "mlx_build_benchmark": 6, "mlx_build_python_bind": 6, "multipl": [6, 118, 125, 134, 135, 242, 251, 285, 288], "wish": 6, "variabl": [6, 98, 103, 183, 185, 186], "export": 6, "developer_dir": 6, "app": 6, "content": [6, 228], "sdk": 6, "xcrun": 6, "macosx": 6, "show": [6, 205], "unabl": 6, "tool": 6, "select": [6, 187, 224, 228], "sudo": 6, "ouptut": 6, "finder": 6, "iterm": 6, "termin": 6, "click": 6, "uncheck": 6, "window": 6, "restart": 6, "devicetyp": 7, "attribut": [7, 8, 9, 26, 278], "kwarg": [8, 154, 155, 290], "union": [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 28, 29, 30, 31, 32, 33, 34, 36, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 51, 53, 54, 55, 56, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 72, 73, 74, 75, 76, 77, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 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, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 140, 141, 143, 144, 145, 146, 147, 148, 149, 150, 158, 159, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 187, 188, 189, 213, 229, 231, 239], "absolut": [10, 13, 255, 256, 267], "semant": [11, 61, 73, 74, 75, 99, 100, 104, 105, 113, 118, 120, 123, 125, 171, 289], "keepdim": [12, 14, 22, 23, 29, 30, 31, 32, 40, 41, 42, 43, 45, 56, 59, 106, 117, 119, 121, 122, 133, 163, 172, 184], "reduct": [12, 14, 117, 119, 122, 133, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268], "reduc": [12, 14, 22, 23, 117, 119, 121, 122, 133, 160, 172, 184, 211, 264], "unspecifi": [12, 14, 15, 22, 23, 24, 25, 64, 97, 117, 119, 121, 122, 128, 132, 133, 147, 163, 164, 172, 174, 184, 188, 290], "entir": [12, 14, 22, 23, 117, 119, 121, 122, 133, 172, 184, 215, 216], "singleton": [12, 14, 22, 23, 117, 118, 119, 121, 122, 133, 172, 184], "rtol": 13, "05": [13, 211, 219, 220, 221, 245], "atol": 13, "08": [13, 192, 193, 194, 195, 199, 258], "approxim": [13, 160, 218, 254, 255, 256], "comparison": [13, 75, 99, 100, 104, 105], "equal": [13, 24, 60, 81, 100, 105, 132, 141, 165, 220, 222], "ab": [13, 106, 183, 219, 220, 221, 223, 245, 269], "array_equ": 13, "rel": 13, "toler": 13, "boolean": [13, 60, 114, 115, 116, 205, 238, 284], "interv": [15, 107, 141, 145], "increment": 15, "otherwis": [15, 201, 202, 229, 231, 239, 252, 253, 261, 267, 274, 285, 286], "int32": [15, 106, 141, 205, 284, 287], "convent": [15, 67, 194], "lead": 15, "fraction": 15, "integr": [15, 174, 285], "invers": [16, 17, 18, 19, 20, 21, 77, 85, 86, 87, 88, 89, 90], "cosin": [16, 17, 68, 69, 247, 258], "hyperbol": [17, 19, 21, 69, 162, 177], "sine": [18, 19, 161, 162, 247], "minimum": [23, 63, 251, 258], "kth": [24, 132], "partit": 24, "order": [24, 106, 132, 134, 209, 219, 240, 249], "undefin": [24, 132, 284], "sort": [24, 25, 132], "flatten": [24, 25, 106, 130, 132, 147, 164, 174, 175, 201], "dimension": [26, 82, 83, 84, 85, 86, 87, 91, 92, 93, 211, 212, 213, 217, 222, 244, 251, 284, 286], "val": [26, 97], "tupl": [26, 47, 58, 64, 66, 74, 78, 80, 103, 106, 131, 134, 148, 168, 183, 185, 193, 194, 195, 196, 201, 202, 203, 213, 231, 233, 247, 249, 283], "ndarrai": [26, 284, 285, 287], "properti": [27, 35, 44, 50, 52, 238], "argument": [27, 47, 58, 78, 98, 183, 202, 209, 281, 288, 289, 290], "decim": [48, 149], "indices_or_sect": [53, 165], "nest": [57, 209, 278, 283], "ddof": [59, 184], "equal_nan": 60, "nan": 60, "a_min": 63, "a_max": 63, "edg": [63, 131], "At": 63, "anoth": [63, 118, 171, 187, 209, 224, 284, 289], "pad": [65, 66, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 212, 213], "dilat": [65, 66], "group": [65, 66, 72, 134, 135, 219, 244], "1d": [65, 67, 175], "convolut": [65, 66, 67, 212, 213, 215, 216], "channel": [65, 66, 211, 212, 213, 215, 216], "c_in": [65, 66], "c_out": [65, 66], "convolv": [65, 66], "2d": [66, 134, 211, 215], "spatial": [66, 219], "symmetr": 66, "discret": [67, 82, 83, 84, 85, 86, 87, 91, 92, 93, 217], "swap": [67, 173, 241, 244], "conv": 67, "filter": [67, 212, 213, 224, 228], "flip": 67, "signal": 67, "bias": [72, 134, 135, 229, 239, 242], "group_siz": [72, 134, 135, 244], "64": [72, 134, 135, 205, 244], "configur": 72, "formal": [72, 134], "notat": [72, 201, 233], "quantiz": [72, 108, 135, 244], "w_i": [72, 134], "hat": [72, 134], "occupi": [72, 134, 135], "divis": [73, 96, 134], "quotient": [73, 74, 96], "remaind": 74, "fuction": 74, "faster": [74, 254], "mathrm": [76, 158, 220], "frac": [76, 134, 158, 191, 192, 193, 194, 195, 199, 211, 214, 215, 216, 219, 220, 221, 222, 245, 258, 261, 264], "pi": [76, 251], "int_0": 76, "dx": 76, "erf": 77, "node": [78, 160, 186], "dict": [78, 108, 152, 153, 154, 234, 237, 240, 241, 278, 283, 288], "leaf": [78, 201, 202, 228], "exponenti": [79, 248, 272], "insert": [80, 289], "ident": [81, 170, 236], "diagon": [81, 180, 181, 182], "zero": [81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 180, 181, 182, 189, 209, 214, 215, 216, 231, 284], "th": 81, "whose": [81, 190], "One": [82, 85, 91, 150], "fourier": [82, 83, 84, 85, 86, 87, 91, 92, 93], "truncat": [82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 144], "dft": [82, 83, 84, 85, 86, 87, 91, 92, 93], "rfft": 88, "real": [88, 89, 90, 91, 92, 93], "rfft2": 89, "rfftn": 90, "silent": [91, 92, 93], "start_axi": 94, "end_axi": 94, "integ": [96, 106, 131, 134, 135, 136, 141, 165, 178, 186, 205, 217, 284], "floor": 96, "fun": [98, 103, 183, 185, 186, 284, 285, 289], "argnam": [98, 183], "cpp_function": [98, 183, 186], "neither": [98, 183], "keyword": [98, 154, 155, 183, 202, 209, 281, 288, 290], "strict": [99, 104, 229, 231, 239], "ordinari": 102, "ord": 106, "tabl": [106, 205, 217], "frobeniu": 106, "matric": 106, "strictli": 106, "mathemat": 106, "variou": 106, "purpos": 106, "calcul": 106, "fro": 106, "inf": [106, 242], "largest": 106, "sing": 106, "smallest": 106, "singular": 106, "nuclear": 106, "_f": 106, "sum_": [106, 264], "a_": 106, "valueerror": [106, 231], "refer": [106, 220, 223, 269, 284], "golub": 106, "van": 106, "loan": 106, "baltimor": 106, "md": 106, "john": 106, "hopkin": 106, "univers": 106, "1985": 106, "pg": 106, "la": 106, "arang": [106, 284, 286], "9": [106, 191, 193, 194, 195, 196, 286], "74597": 106, "20": 106, "84804": 106, "41421": 106, "23607": 106, "74166": 106, "24264": 106, "11": 106, "225": 106, "50": 107, "evenli": 107, "binari": [108, 151, 152, 153, 154, 155, 252, 257, 274], "npy": [108, 151, 288], "safetensor": [108, 153, 285, 288], "gguf": [108, 152, 288], "unsupport": 108, "tensor": [108, 178, 268, 286], "natur": [109, 111, 285], "logarithm": [109, 110, 111, 112], "log": [111, 113, 117, 262, 264, 266], "plu": 111, "exp": [113, 117, 138, 163, 248, 262, 272, 289], "stabl": [113, 117, 163, 264], "prepend": 118, "remov": [118, 137, 168], "negat": 126, "beforehand": 130, "pad_with": 131, "constant_valu": 131, "pad_width": 131, "before_1": 131, "after_1": 131, "before_2": 131, "after_2": 131, "before_n": 131, "after_n": 131, "before_i": 131, "after_i": 131, "extend": 131, "side": 131, "smaller": [132, 196], "everi": [134, 160, 202], "particular": [134, 219], "consecut": [134, 247], "w_1": 134, "w_g": 134, "begin": [134, 248, 252, 261, 267, 272, 274], "align": 134, "max_i": 134, "min_i": 134, "textrm": [134, 218, 254], "round": 134, "pack": [134, 135], "unsign": [134, 135, 205], "lower": [134, 141, 144, 145, 180], "upper": [134, 141, 144, 145], "1st": 134, "signific": 134, "2nd": 134, "dequant": 134, "w_q": 134, "whether": [135, 228, 242], "prng": [136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 281], "num_sampl": 137, "unnorm": [137, 257, 259], "draw": 137, "uint32": [137, 205], "cdf": [138, 218, 254], "accord": [138, 187, 242], "seed": 139, "low": [141, 145], "high": [141, 145, 209, 217], "bound": [141, 144, 145, 218, 284, 289], "roadcast": 141, "domain": 144, "uniformli": 145, "repetit": 147, "preserv": 148, "reciproc": 150, "arr": [151, 284], "uncompress": 154, "my_path": 154, "tree_flatten": [154, 202, 203, 209], "transformerencod": 154, "128": [154, 209], "flat_param": 154, "compress": 155, "simplif": 160, "reus": 160, "consumpt": 160, "meant": 160, "overhead": [160, 285, 289], "1m": 160, "thousand": [160, 285], "foo": 160, "matmul": [160, 289], "twice": [160, 289], "subarrai": 165, "being": [170, 209], "prevent": [170, 268, 286], "flow": [170, 285], "unchang": [170, 247], "axis1": 173, "axis2": 173, "taken": 174, "prior": [174, 175], "exclud": 175, "dot": [178, 201, 233, 242], "elsewher": [180, 284], "col": 180, "triangl": 180, "mse": 183, "param": [183, 209], "lvalu": 183, "dlvalu": 183, "dparam": 183, "lasso": 183, "l1": [183, 261, 263, 264, 267], "varianc": [184, 211, 219], "divisor": 184, "cotang": 185, "in_ax": 186, "out_ax": 186, "prefix": [186, 201], "fn": [190, 202, 287], "callabl": [190, 201, 202, 224, 225, 228, 249, 253], "wrt": 190, "rho": 191, "06": [191, 268], "paper": [191, 192, 193, 195, 196, 211, 251], "zeiler": 191, "2012": [191, 199], "adapt": [191, 192], "1212": 191, "5701": 191, "v_": [191, 192, 193, 194, 195, 199, 200], "v_t": [191, 192, 193, 194, 195, 199, 200], "g_t": [191, 192, 193, 194, 195, 196, 199, 200], "delta": [191, 261], "w_": [191, 192, 193, 194, 195, 196, 199, 200], "u_t": 191, "epsilon": [191, 192, 193, 194, 195, 199, 211, 219, 220, 221, 245, 258], "u_": 191, "w_t": [191, 192, 193, 194, 195, 196, 199, 200], "lambda": [191, 192, 193, 194, 195, 196, 199, 200, 202, 209, 224, 229, 248, 272], "averag": [191, 193, 194, 195], "denomin": [191, 192, 193, 194, 195, 199, 220, 258], "stabil": [191, 192, 193, 194, 195, 199, 211, 219, 220, 221, 245, 258], "duchi": 192, "hazan": 192, "singer": 192, "2011": 192, "subgradi": 192, "onlin": 192, "stochast": [192, 193, 195, 200, 285], "jmlr": 192, "999": [193, 194, 195], "omit": [193, 195], "estim": [193, 195], "kingma": [193, 195], "ba": [193, 195], "2015": [193, 195, 215], "iclr": [193, 194, 195], "m_": [193, 194, 195, 196], "beta_1": [193, 194, 195, 196], "m_t": [193, 194, 195, 196], "beta_2": [193, 194, 195, 196], "weight_decai": [194, 196, 200], "contrast": [194, 198], "loshchilov": 194, "hutter": 194, "decoupl": 194, "decai": [194, 196, 200], "regular": [194, 215, 223, 269, 284], "adam": [195, 196], "infin": 195, "99": [196, 199], "sign": [196, 205], "tend": 196, "larger": [196, 247], "10x": 196, "adamw": 196, "maintain": [196, 215, 216], "strength": [196, 200], "wd": 196, "chen": 196, "symbol": 196, "discoveri": 196, "2302": 196, "06675": 196, "c_": 196, "eta": 196, "c_t": 196, "momentum": [196, 200, 211], "basi": 197, "appli": [197, 202, 209, 211, 212, 213, 215, 216, 218, 219, 220, 221, 222, 223, 225, 236, 243, 244, 245, 246, 248, 250, 252, 254, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274], "optimizerst": 197, "recurs": [198, 209, 228, 229, 234, 237, 239, 278], "defaultdict": 198, "miss": [198, 231, 288], "present": 198, "tieleman": 199, "hinton": 199, "lectur": 199, "coursera": 199, "smooth": [199, 259, 267], "dampen": 200, "nesterov": 200, "descent": [200, 285], "mu": 200, "tau": 200, "l2": [200, 261, 264], "penalti": 200, "is_leaf": [201, 202], "arbitrari": [201, 278], "depth": [201, 216], "hello": [201, 203], "charact": 201, "flat": [201, 203], "superset": 202, "extra": 202, "closer": 202, "constitut": 202, "dict_kei": 202, "recreat": 203, "world": 203, "42": 203, "byte": 205, "bool_": 205, "uint8": 205, "uint16": 205, "16": [205, 220, 224, 278], "uint64": 205, "int8": 205, "int16": 205, "int64": 205, "arbitrarili": [209, 283, 287], "done": [209, 214, 285, 286], "manual": 209, "explicitli": [209, 281], "solv": 209, "intuit": 209, "freez": [209, 239, 278], "finetun": 209, "in_dim": [209, 278], "out_dim": [209, 278], "enumer": 209, "caus": [209, 285], "local": [209, 215], "scope": 209, "l2_loss": 209, "y_hat": 209, "trainable_paramet": [209, 228], "loss_and_grad": 209, "workhors": 209, "Its": 209, "frozen": [209, 229, 237, 239, 244, 278], "individu": [209, 215, 216], "subset": [209, 228], "action": 209, "displai": 209, "tree_map": 209, "count": 209, "num_param": 209, "preclud": 209, "pure": [209, 280], "pattern": [209, 285], "achiev": 209, "other_input": 209, "necessari": 209, "wrap": 209, "apply_to_modul": [209, 229], "children": 209, "filter_and_map": 209, "leaf_modul": 209, "load_weight": [209, 285], "named_modul": 209, "save_weight": 209, "unfreez": [209, 229], "update_modul": 209, "sequenti": 209, "relu": [209, 243, 253, 270], "prelu": 209, "gelu": [209, 255, 256], "silu": 209, "selu": 209, "mish": 209, "quantizedlinear": 209, "conv1d": 209, "conv2d": 209, "batchnorm": 209, "layernorm": 209, "groupnorm": 209, "instancenorm": 209, "dropout": [209, 215, 216, 236, 253], "dropout2d": 209, "dropout3d": 209, "alibi": 209, "sinusoidalpositionalencod": 209, "gelu_approx": [209, 218, 254], "gelu_fast_approx": [209, 218, 254], "binary_cross_entropi": 209, "kl_div_loss": 209, "l1_loss": 209, "mse_loss": 209, "nll_loss": 209, "smooth_l1_loss": 209, "triplet_loss": 209, "hinge_loss": 209, "huber_loss": 209, "log_cosh_loss": 209, "cosine_similarity_loss": 209, "affin": [211, 219, 220, 221, 222, 244], "track_running_stat": 211, "var": [211, 219, 220, 221], "gamma": [211, 219, 220, 221, 245], "nc": 211, "nlc": [211, 212], "four": 211, "nhwc": [211, 213], "height": [211, 213, 215, 216], "width": [211, 213, 215, 216, 244], "deep": 211, "intern": 211, "covari": 211, "shift": 211, "bn": 211, "in_channel": [212, 213], "out_channel": [212, 213], "kernel_s": [212, 213], "learnabl": [212, 213, 249], "portion": 214, "dure": [214, 215, 216, 286], "independ": [215, 216], "nwhc": 215, "whc": 215, "entri": [215, 216], "benefici": [215, 216, 285], "earli": 215, "adjac": 215, "pixel": 215, "correl": 215, "thompson": 215, "goroshin": 215, "jain": 215, "lecun": 215, "bregler": 215, "cvpr": 215, "ndhwc": 216, "dhwc": 216, "medic": 216, "video": 216, "num_embed": 217, "lookup": 217, "typic": [217, 280, 285], "usual": [217, 283, 285], "vocabulari": 217, "approx": 218, "unit": [218, 246, 248, 250, 254, 255, 256, 271, 272, 273], "phi": [218, 254], "geluapprox": 218, "sigma": [218, 250, 255, 256, 273], "60033": [218, 255], "0433603": [218, 255], "gelufast": 218, "773": [218, 256], "regard": 218, "num_group": 219, "pytorch_compat": 219, "split": 219, "preced": 219, "http": [219, 220, 221, 223, 245, 269], "org": [219, 220, 221, 223, 245, 269], "1803": 219, "08494": 219, "inorm": 220, "1607": [220, 221], "08022": 220, "06450": 221, "uniform": [222, 231, 281, 289], "mathcal": 222, "u": 222, "d_i": 222, "monoton": [223, 269], "1908": [223, 269], "08681": [223, 269], "tanh": [223, 269], "softplu": [223, 269], "map_fn": [224, 228], "filter_fn": [224, 228], "valid_parameter_filt": 224, "apply_fn": 225, "descend": 226, "is_leaf_fn": 228, "found": 228, "drop": 228, "idempot": [229, 239], "attent": [229, 242, 251, 253], "endswith": 229, "file_or_weight": 231, "ok": 231, "certain": 236, "ie": 239, "noop": 239, "unfrozen": 239, "chang": [240, 244, 261, 267, 286], "tracer": 240, "partial": [240, 241, 285], "child": 241, "programmat": 241, "query_input_dim": 242, "key_input_dim": 242, "value_input_dim": 242, "value_dim": 242, "value_output_dim": 242, "head": [242, 253], "aggreg": 242, "linearli": 242, "neg": [242, 266, 268, 284], "attend": 242, "num_paramet": 243, "init": 243, "25": 243, "parametr": [243, 270], "classmethod": 244, "from_linear": 244, "quantize_modul": 244, "1910": 245, "07467": 245, "rectifi": [246, 271], "10000": 247, "rotat": 247, "slightli": [247, 289], "angular": 247, "frequenc": [247, 251], "_cos_sin_theta_kei": 247, "precomput": 247, "_cos_sin_theta_valu": 247, "leq": [248, 261, 272], "0507": [248, 272], "67326": [248, 272], "elu": [248, 272], "plain": 249, "known": [250, 273], "swish": [250, 273], "cdot": [250, 255, 256, 258, 260, 273], "min_freq": 251, "0001": 251, "max_freq": 251, "cos_first": 251, "full_turn": 251, "sinusoid": 251, "sin": [251, 287], "threshold": [252, 261, 267, 274], "geq": [252, 274], "num_encoder_lay": 253, "num_decoder_lay": 253, "custom_encod": 253, "custom_decod": 253, "norm_first": 253, "decod": 253, "interact": 253, "mechan": 253, "hidden": 253, "exact": [255, 256], "0003": 255, "015": 256, "pre": 257, "predict": [257, 259, 260, 261, 262, 263, 264, 265, 266, 267], "105361": 257, "223144": 257, "20397": 257, "916291": 257, "612192": 257, "x1": 258, "x2": 258, "x_1": 258, "x_2": 258, "label_smooth": 259, "hing": 260, "y_": [260, 264], "pred": [260, 264], "huber": 261, "l_": [261, 268], "kullback": 262, "leibler": 262, "diverg": 262, "cosh": 264, "logcosh": 264, "sensit": 264, "outlier": 264, "dual": 264, "behavior": [264, 284, 285], "offer": 264, "balanc": 264, "robust": 264, "approach": 264, "task": 264, "likelihood": 266, "nll": 266, "formula": 267, "anchor": 268, "margin": 268, "triplet": 268, "_p": 268, "degre": 268, "pairwis": 268, "instabl": 268, "subclass": 278, "concept": 278, "mymlp": 278, "in_proj": 278, "subsequ": 280, "implicit": 281, "fine": [281, 285], "grain": 281, "control": [281, 285], "manag": [281, 289], "pseudo": 281, "altern": 281, "splittabl": 281, "threefri": 281, "counter": 281, "cycl": 283, "slice": 284, "ellipsi": 284, "syntax": 284, "idx": 284, "mix": 284, "take_along_axi": 284, "lack": 284, "propag": 284, "extrem": [284, 285], "ineffici": [284, 285], "nonzero": 284, "reflect": [284, 286], "dfdx": 284, "record": 285, "nice": 285, "rerun": 285, "dynam": 285, "easier": 285, "worri": 285, "fun1": 285, "expensive_fun": 285, "cost": 285, "code": 285, "consum": 285, "eager": 285, "thank": 285, "weights_fp16": 285, "trade": 285, "too": 285, "bad": 285, "idea": 285, "On": 285, "grow": 285, "computation": 285, "costli": 285, "wide": 285, "pretti": 285, "ten": 285, "okai": 285, "outer": 285, "value_and_grad_fn": 285, "awar": 285, "implicitli": 285, "anytim": 285, "memoryview": [285, 286], "perfectli": 285, "first_lay": 285, "second_layer_a": 285, "second_layer_b": 285, "frequent": 285, "protocol": 286, "receiv": 286, "pep": 286, "3118": 286, "view": 286, "a_view": 286, "owndata": 286, "quit": 286, "power": 286, "extern": 286, "x_view": 286, "modifi": 286, "df": 286, "x\u00b2": 286, "2x": 286, "indirectli": 286, "modif": 286, "seen": 286, "occur": 286, "incorpor": 286, "issu": 286, "incorrect": 286, "experiment": 286, "break": 286, "advis": 286, "intermedi": 286, "jnp": 286, "tf": 286, "inspect": 287, "page": 287, "composit": 287, "archiv": 288, "savez_compress": 288, "save_safetensor": 288, "save_gguf": 288, "arr_0": 288, "pool": 289, "advantag": 289, "don": 289, "parallel": 289, "race": 289, "interest": 289, "albeit": 289, "contriv": 289, "suppos": 289, "d1": 289, "d2": 289, "4096": 289, "dens": 289, "better": 289, "millisecond": 289, "measur": 289, "default_stream": 290, "default_devic": 290, "my_devic": 290}, "objects": {"mlx.core": [[7, 0, 1, "", "Device"], [8, 0, 1, "", "Dtype"], [9, 0, 1, "", "Stream"], [10, 2, 1, "", "abs"], [11, 2, 1, "", "add"], [12, 2, 1, "", "all"], [13, 2, 1, "", "allclose"], [14, 2, 1, "", "any"], [15, 2, 1, "", "arange"], [16, 2, 1, "", "arccos"], [17, 2, 1, "", "arccosh"], [18, 2, 1, "", "arcsin"], [19, 2, 1, "", "arcsinh"], [20, 2, 1, "", "arctan"], [21, 2, 1, "", "arctanh"], [22, 2, 1, "", "argmax"], [23, 2, 1, "", "argmin"], [24, 2, 1, "", "argpartition"], [25, 2, 1, "", "argsort"], [26, 0, 1, "", "array"], [60, 2, 1, "", "array_equal"], [61, 2, 1, "", "broadcast_to"], [62, 2, 1, "", "ceil"], [63, 2, 1, "", "clip"], [64, 2, 1, "", "concatenate"], [65, 2, 1, "", "conv1d"], [66, 2, 1, "", "conv2d"], [67, 2, 1, "", "convolve"], [68, 2, 1, "", "cos"], [69, 2, 1, "", "cosh"], [70, 2, 1, "", "default_device"], [71, 2, 1, "", "default_stream"], [72, 2, 1, "", "dequantize"], [73, 2, 1, "", "divide"], [74, 2, 1, "", "divmod"], [75, 2, 1, "", "equal"], [76, 2, 1, "", "erf"], [77, 2, 1, "", "erfinv"], [78, 2, 1, "", "eval"], [79, 2, 1, "", "exp"], [80, 2, 1, "", "expand_dims"], [81, 2, 1, "", "eye"], [94, 2, 1, "", "flatten"], [95, 2, 1, "", "floor"], [96, 2, 1, "", "floor_divide"], [97, 2, 1, "", "full"], [98, 2, 1, "", "grad"], [99, 2, 1, "", "greater"], [100, 2, 1, "", "greater_equal"], [101, 2, 1, "", "identity"], [102, 2, 1, "", "inner"], [103, 2, 1, "", "jvp"], [104, 2, 1, "", "less"], [105, 2, 1, "", "less_equal"], [107, 2, 1, "", "linspace"], [108, 2, 1, "", "load"], [109, 2, 1, "", "log"], [110, 2, 1, "", "log10"], [111, 2, 1, "", "log1p"], [112, 2, 1, "", "log2"], [113, 2, 1, "", "logaddexp"], [114, 2, 1, "", "logical_and"], [115, 2, 1, "", "logical_not"], [116, 2, 1, "", "logical_or"], [117, 2, 1, "", "logsumexp"], [118, 2, 1, "", "matmul"], [119, 2, 1, "", "max"], [120, 2, 1, "", "maximum"], [121, 2, 1, "", "mean"], [122, 2, 1, "", "min"], [123, 2, 1, "", "minimum"], [124, 2, 1, "", "moveaxis"], [125, 2, 1, "", "multiply"], [126, 2, 1, "", "negative"], [127, 2, 1, "", "new_stream"], [128, 2, 1, "", "ones"], [129, 2, 1, "", "ones_like"], [130, 2, 1, "", "outer"], [131, 2, 1, "", "pad"], [132, 2, 1, "", "partition"], [133, 2, 1, "", "prod"], [134, 2, 1, "", "quantize"], [135, 2, 1, "", "quantized_matmul"], [146, 2, 1, "", "reciprocal"], [147, 2, 1, "", "repeat"], [148, 2, 1, "", "reshape"], [149, 2, 1, "", "round"], [150, 2, 1, "", "rsqrt"], [151, 2, 1, "", "save"], [152, 2, 1, "", "save_gguf"], [153, 2, 1, "", "save_safetensors"], [154, 2, 1, "", "savez"], [155, 2, 1, "", "savez_compressed"], [156, 2, 1, "", "set_default_device"], [157, 2, 1, "", "set_default_stream"], [158, 2, 1, "", "sigmoid"], [159, 2, 1, "", "sign"], [160, 2, 1, "", "simplify"], [161, 2, 1, "", "sin"], [162, 2, 1, "", "sinh"], [163, 2, 1, "", "softmax"], [164, 2, 1, "", "sort"], [165, 2, 1, "", "split"], [166, 2, 1, "", "sqrt"], [167, 2, 1, "", "square"], [168, 2, 1, "", "squeeze"], [169, 2, 1, "", "stack"], [170, 2, 1, "", "stop_gradient"], [171, 2, 1, "", "subtract"], [172, 2, 1, "", "sum"], [173, 2, 1, "", "swapaxes"], [174, 2, 1, "", "take"], [175, 2, 1, "", "take_along_axis"], [176, 2, 1, "", "tan"], [177, 2, 1, "", "tanh"], [178, 2, 1, "", "tensordot"], [179, 2, 1, "", "transpose"], [180, 2, 1, "", "tri"], [181, 2, 1, "", "tril"], [182, 2, 1, "", "triu"], [183, 2, 1, "", "value_and_grad"], [184, 2, 1, "", "var"], [185, 2, 1, "", "vjp"], [186, 2, 1, "", "vmap"], [187, 2, 1, "", "where"], [188, 2, 1, "", "zeros"], [189, 2, 1, "", "zeros_like"]], "mlx.core.Device": [[7, 1, 1, "", "__init__"]], "mlx.core.Dtype": [[8, 1, 1, "", "__init__"]], "mlx.core.Stream": [[9, 1, 1, "", "__init__"]], "mlx.core.array": [[27, 3, 1, "", "T"], [26, 1, 1, "", "__init__"], [28, 1, 1, "", "abs"], [29, 1, 1, "", "all"], [30, 1, 1, "", "any"], [31, 1, 1, "", "argmax"], [32, 1, 1, "", "argmin"], [33, 1, 1, "", "astype"], [34, 1, 1, "", "cos"], [35, 3, 1, "", "dtype"], [36, 1, 1, "", "exp"], [37, 1, 1, "", "item"], [38, 1, 1, "", "log"], [39, 1, 1, "", "log1p"], [40, 1, 1, "", "logsumexp"], [41, 1, 1, "", "max"], [42, 1, 1, "", "mean"], [43, 1, 1, "", "min"], [44, 3, 1, "", "ndim"], [45, 1, 1, "", "prod"], [46, 1, 1, "", "reciprocal"], [47, 1, 1, "", "reshape"], [48, 1, 1, "", "round"], [49, 1, 1, "", "rsqrt"], [50, 3, 1, "", "shape"], [51, 1, 1, "", "sin"], [52, 3, 1, "", "size"], [53, 1, 1, "", "split"], [54, 1, 1, "", "sqrt"], [55, 1, 1, "", "square"], [56, 1, 1, "", "sum"], [57, 1, 1, "", "tolist"], [58, 1, 1, "", "transpose"], [59, 1, 1, "", "var"]], "mlx.core.fft": [[82, 2, 1, "", "fft"], [83, 2, 1, "", "fft2"], [84, 2, 1, "", "fftn"], [85, 2, 1, "", "ifft"], [86, 2, 1, "", "ifft2"], [87, 2, 1, "", "ifftn"], [88, 2, 1, "", "irfft"], [89, 2, 1, "", "irfft2"], [90, 2, 1, "", "irfftn"], [91, 2, 1, "", "rfft"], [92, 2, 1, "", "rfft2"], [93, 2, 1, "", "rfftn"]], "mlx.core.linalg": [[106, 2, 1, "", "norm"]], "mlx.core.random": [[136, 2, 1, "", "bernoulli"], [137, 2, 1, "", "categorical"], [138, 2, 1, "", "gumbel"], [139, 2, 1, "", "key"], [140, 2, 1, "", "normal"], [141, 2, 1, "", "randint"], [142, 2, 1, "", "seed"], [143, 2, 1, "", "split"], [144, 2, 1, "", "truncated_normal"], [145, 2, 1, "", "uniform"]], "mlx.nn": [[210, 0, 1, "", "ALiBi"], [211, 0, 1, "", "BatchNorm"], [212, 0, 1, "", "Conv1d"], [213, 0, 1, "", "Conv2d"], [214, 0, 1, "", "Dropout"], [215, 0, 1, "", "Dropout2d"], [216, 0, 1, "", "Dropout3d"], [217, 0, 1, "", "Embedding"], [218, 0, 1, "", "GELU"], [219, 0, 1, "", "GroupNorm"], [220, 0, 1, "", "InstanceNorm"], [221, 0, 1, "", "LayerNorm"], [222, 0, 1, "", "Linear"], [223, 0, 1, "", "Mish"], [278, 0, 1, "", "Module"], [242, 0, 1, "", "MultiHeadAttention"], [243, 0, 1, "", "PReLU"], [244, 0, 1, "", "QuantizedLinear"], [245, 0, 1, "", "RMSNorm"], [246, 0, 1, "", "ReLU"], [247, 0, 1, "", "RoPE"], [248, 0, 1, "", "SELU"], [249, 0, 1, "", "Sequential"], [250, 0, 1, "", "SiLU"], [251, 0, 1, "", "SinusoidalPositionalEncoding"], [252, 0, 1, "", "Step"], [253, 0, 1, "", "Transformer"], [254, 0, 1, "", "gelu"], [255, 0, 1, "", "gelu_approx"], [256, 0, 1, "", "gelu_fast_approx"], [269, 0, 1, "", "mish"], [270, 0, 1, "", "prelu"], [271, 0, 1, "", "relu"], [272, 0, 1, "", "selu"], [273, 0, 1, "", "silu"], [274, 0, 1, "", "step"], [190, 2, 1, "", "value_and_grad"]], "mlx.nn.Module": [[224, 1, 1, "", "apply"], [225, 1, 1, "", "apply_to_modules"], [226, 1, 1, "", "children"], [227, 1, 1, "", "eval"], [228, 1, 1, "", "filter_and_map"], [229, 1, 1, "", "freeze"], [230, 1, 1, "", "leaf_modules"], [231, 1, 1, "", "load_weights"], [232, 1, 1, "", "modules"], [233, 1, 1, "", "named_modules"], [234, 1, 1, "", "parameters"], [235, 1, 1, "", "save_weights"], [236, 1, 1, "", "train"], [237, 1, 1, "", "trainable_parameters"], [238, 3, 1, "", "training"], [239, 1, 1, "", "unfreeze"], [240, 1, 1, "", "update"], [241, 1, 1, "", "update_modules"]], "mlx.nn.RoPE": [[247, 4, 1, "", "_cos_sin_theta_key"], [247, 4, 1, "", "_cos_sin_theta_value"]], "mlx.nn.losses": [[257, 0, 1, "", "binary_cross_entropy"], [258, 0, 1, "", "cosine_similarity_loss"], [259, 0, 1, "", "cross_entropy"], [260, 0, 1, "", "hinge_loss"], [261, 0, 1, "", "huber_loss"], [262, 0, 1, "", "kl_div_loss"], [263, 0, 1, "", "l1_loss"], [264, 0, 1, "", "log_cosh_loss"], [265, 0, 1, "", "mse_loss"], [266, 0, 1, "", "nll_loss"], [267, 0, 1, "", "smooth_l1_loss"], [268, 0, 1, "", "triplet_loss"]], "mlx.optimizers": [[191, 0, 1, "", "AdaDelta"], [192, 0, 1, "", "Adagrad"], [193, 0, 1, "", "Adam"], [194, 0, 1, "", "AdamW"], [195, 0, 1, "", "Adamax"], [196, 0, 1, "", "Lion"], [197, 0, 1, "", "Optimizer"], [198, 0, 1, "", "OptimizerState"], [199, 0, 1, "", "RMSprop"], [200, 0, 1, "", "SGD"]], "mlx.optimizers.Optimizer": [[197, 4, 1, "", "state"]], "mlx.utils": [[201, 2, 1, "", "tree_flatten"], [202, 2, 1, "", "tree_map"], [203, 2, 1, "", "tree_unflatten"]]}, "objtypes": {"0": "py:class", "1": "py:method", "2": "py:function", "3": "py:property", "4": "py:attribute"}, "objnames": {"0": ["py", "class", "Python class"], "1": ["py", "method", "Python method"], "2": ["py", "function", "Python function"], "3": ["py", "property", "Python property"], "4": ["py", "attribute", "Python attribute"]}, "titleterms": {"oper": [0, 1, 279], "develop": 1, "document": 1, "introduc": 1, "exampl": [1, 5, 289], "primit": 1, "us": [1, 285, 290], "implement": [1, 3], "cpu": 1, "backend": 1, "gpu": 1, "transform": [1, 253, 282, 285, 287], "build": [1, 6], "bind": 1, "python": [1, 5, 6], "cmake": 1, "setuptool": 1, "usag": [1, 5], "result": 1, "script": [1, 3], "download": [1, 3], "code": [1, 3], "linear": [2, 208, 222], "regress": 2, "llm": 3, "infer": 3, "model": 3, "attent": 3, "layer": [3, 4, 276], "encod": 3, "full": [3, 97], "gener": 3, "put": 3, "all": [3, 12, 29], "togeth": 3, "convert": 3, "weight": 3, "load": [3, 108, 288], "benchmark": 3, "multi": 4, "perceptron": 4, "mlx": [5, 7, 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, 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, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274], "instal": [5, 6], "api": [5, 6], "refer": 5, "c": [5, 6], "further": 5, "read": 5, "from": [6, 284], "pypi": 6, "troubleshoot": 6, "sourc": 6, "requir": 6, "option": 6, "metal": 6, "found": 6, "x86": 6, "shell": 6, "core": [7, 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], "devic": [7, 206], "dtype": [8, 35], "stream": [9, 206, 290], "ab": [10, 28], "add": 11, "allclos": 13, "ani": [14, 30], "arang": 15, "arcco": 16, "arccosh": 17, "arcsin": 18, "arcsinh": 19, "arctan": 20, "arctanh": 21, "argmax": [22, 31], "argmin": [23, 32], "argpartit": 24, "argsort": 25, "arrai": [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, 204, 284, 288], "t": 27, "astyp": 33, "co": [34, 68], "exp": [36, 79], "item": 37, "log": [38, 109], "log1p": [39, 111], "logsumexp": [40, 117], "max": [41, 119], "mean": [42, 121], "min": [43, 122], "ndim": 44, "prod": [45, 133], "reciproc": [46, 146], "reshap": [47, 148], "round": [48, 149], "rsqrt": [49, 150], "shape": 50, "sin": [51, 161], "size": 52, "split": [53, 143, 165], "sqrt": [54, 166], "squar": [55, 167], "sum": [56, 172], "tolist": 57, "transpos": [58, 179], "var": [59, 184], "array_equ": 60, "broadcast_to": 61, "ceil": 62, "clip": 63, "concaten": 64, "conv1d": [65, 212], "conv2d": [66, 213], "convolv": 67, "cosh": 69, "default_devic": 70, "default_stream": 71, "dequant": 72, "divid": 73, "divmod": 74, "equal": 75, "erf": 76, "erfinv": 77, "eval": [78, 227], "expand_dim": 80, "ey": 81, "fft": [82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 207], "fft2": 83, "fftn": 84, "ifft": 85, "ifft2": 86, "ifftn": 87, "irfft": 88, "irfft2": 89, "irfftn": 90, "rfft": 91, "rfft2": 92, "rfftn": 93, "flatten": 94, "floor": 95, "floor_divid": 96, "grad": [98, 209], "greater": 99, "greater_equ": 100, "ident": 101, "inner": 102, "jvp": 103, "less": 104, "less_equ": 105, "linalg": 106, "norm": 106, "linspac": 107, "log10": 110, "log2": 112, "logaddexp": 113, "logical_and": 114, "logical_not": 115, "logical_or": 116, "matmul": 118, "maximum": 120, "minimum": 123, "moveaxi": 124, "multipli": 125, "neg": 126, "new_stream": 127, "ones": 128, "ones_lik": 129, "outer": 130, "pad": 131, "partit": 132, "quantiz": 134, "quantized_matmul": 135, "random": [136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 281], "bernoulli": 136, "categor": 137, "gumbel": 138, "kei": 139, "normal": 140, "randint": 141, "seed": 142, "truncated_norm": 144, "uniform": 145, "repeat": 147, "save": [151, 288], "save_gguf": 152, "save_safetensor": 153, "savez": 154, "savez_compress": 155, "set_default_devic": 156, "set_default_stream": 157, "sigmoid": 158, "sign": 159, "simplifi": 160, "sinh": 162, "softmax": 163, "sort": 164, "squeez": 168, "stack": 169, "stop_gradi": 170, "subtract": 171, "swapax": 173, "take": 174, "take_along_axi": 175, "tan": 176, "tanh": 177, "tensordot": 178, "tri": 180, "tril": 181, "triu": 182, "value_and_grad": [183, 190], "vjp": 185, "vmap": 186, "where": 187, "zero": 188, "zeros_lik": 189, "nn": [190, 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, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274], "optim": [191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 280], "adadelta": 191, "adagrad": 192, "adam": 193, "adamw": 194, "adamax": 195, "lion": 196, "optimizerst": 198, "rmsprop": 199, "sgd": 200, "util": [201, 202, 203, 283], "tree_flatten": 201, "tree_map": 202, "tree_unflatten": 203, "data": 205, "type": 205, "support": 205, "algebra": 208, "neural": 209, "network": 209, "quick": [209, 287], "start": [209, 287], "The": 209, "modul": [209, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 278], "class": 209, "paramet": [209, 234], "updat": [209, 240, 284], "inspect": 209, "valu": 209, "alibi": 210, "batchnorm": 211, "dropout": 214, "dropout2d": 215, "dropout3d": 216, "embed": 217, "gelu": [218, 254], "groupnorm": 219, "instancenorm": 220, "layernorm": 221, "mish": [223, 269], "appli": 224, "apply_to_modul": 225, "children": 226, "filter_and_map": 228, "freez": 229, "leaf_modul": 230, "load_weight": 231, "named_modul": 233, "save_weight": 235, "train": [236, 238], "trainable_paramet": 237, "unfreez": 239, "update_modul": 241, "multiheadattent": 242, "prelu": [243, 270], "quantizedlinear": 244, "rmsnorm": 245, "relu": [246, 271], "rope": 247, "selu": [248, 272], "sequenti": 249, "silu": [250, 273], "sinusoidalpositionalencod": 251, "step": [252, 274], "gelu_approx": 255, "gelu_fast_approx": 256, "loss": [257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 277], "binary_cross_entropi": 257, "cosine_similarity_loss": 258, "cross_entropi": 259, "hinge_loss": 260, "huber_loss": 261, "kl_div_loss": 262, "l1_loss": 263, "log_cosh_loss": 264, "mse_loss": 265, "nll_loss": 266, "smooth_l1_loss": 267, "triplet_loss": 268, "function": [275, 277, 287], "tree": 283, "index": 284, "differ": 284, "numpi": [284, 286], "In": 284, "place": 284, "lazi": 285, "evalu": 285, "why": 285, "comput": 285, "graph": [285, 287], "onli": 285, "what": 285, "you": 285, "when": 285, "convers": 286, "other": 286, "framework": 286, "pytorch": 286, "jax": 286, "tensorflow": 286, "guid": 287, "basic": 287, "serial": 288, "format": 288, "unifi": 289, "memori": 289, "A": 289, "simpl": 289, "specifi": 290}, "envversion": {"sphinx.domains.c": 2, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 6, "sphinx.domains.index": 1, "sphinx.domains.javascript": 2, "sphinx.domains.math": 2, "sphinx.domains.python": 3, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.intersphinx": 1, "sphinx": 56}})