linalg.h

Go to the documentation of this file.

// Copyright © 2023 Apple Inc.
 
#pragma once
 
#include <optional>
 
#include "mlx/array.h"
#include "mlx/device.h"
#include "mlx/ops.h"
#include "mlx/stream.h"
 
namespace mlx::core::linalg {

array norm(
    const array& a,
    const double ord,
    const std::optional<std::vector<int>>& axis = std::nullopt,
    bool keepdims = false,
    StreamOrDevice s = {});
inline array norm(
    const array& a,
    const double ord,
    int axis,
    bool keepdims = false,
    StreamOrDevice s = {}) {
  return norm(a, ord, std::vector<int>{axis}, keepdims, s);
}
array norm(
    const array& a,
    const std::string& ord,
    const std::optional<std::vector<int>>& axis = std::nullopt,
    bool keepdims = false,
    StreamOrDevice s = {});
inline array norm(
    const array& a,
    const std::string& ord,
    int axis,
    bool keepdims = false,
    StreamOrDevice s = {}) {
  return norm(a, ord, std::vector<int>{axis}, keepdims, s);
}
array norm(
    const array& a,
    const std::optional<std::vector<int>>& axis = std::nullopt,
    bool keepdims = false,
    StreamOrDevice s = {});
inline array
norm(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}) {
  return norm(a, std::vector<int>{axis}, keepdims, s);
}
 
std::pair<array, array> qr(const array& a, StreamOrDevice s = {});
 
std::vector<array> svd(const array& a, StreamOrDevice s = {});
 
array inv(const array& a, StreamOrDevice s = {});
 
array tri_inv(const array& a, bool upper = false, StreamOrDevice s = {});
 
array cholesky(const array& a, bool upper = false, StreamOrDevice s = {});
 
array pinv(const array& a, StreamOrDevice s = {});
 
array cholesky_inv(const array& a, bool upper = false, StreamOrDevice s = {});

array cross(
    const array& a,
    const array& b,
    int axis = -1,
    StreamOrDevice s = {});
 
array eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 
std::pair<array, array>
eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 
} // namespace mlx::core::linalg