MLX
Loading...
Searching...
No Matches
Classes | Namespaces | Enumerations | Functions
reduce.h File Reference
#include "mlx/backend/common/utils.h"

Go to the source code of this file.

Classes

struct  mlx::core::ReductionPlan
 
struct  mlx::core::DefaultStridedReduce< T, U, Op >
 
struct  mlx::core::DefaultContiguousReduce< T, U, Op >
 

Namespaces

namespace  mlx
 
namespace  mlx::core
 

Enumerations

enum  mlx::core::ReductionOpType {
  mlx::core::ContiguousAllReduce , mlx::core::ContiguousReduce , mlx::core::ContiguousStridedReduce , mlx::core::GeneralContiguousReduce ,
  mlx::core::GeneralStridedReduce , mlx::core::GeneralReduce
}
 

Functions

ReductionPlan mlx::core::get_reduction_plan (const array &x, const std::vector< int > axes)
 
void mlx::core::nd_loop (std::function< void(int)> callback, const std::vector< int > &shape, const std::vector< size_t > &strides)
 
std::pair< std::vector< int >, std::vector< size_t > > mlx::core::shapes_without_reduction_axes (const array &x, const std::vector< int > &axes)
 
template<typename T , typename U , typename OpS , typename OpC , typename Op >
void mlx::core::reduction_op (const array &x, array &out, const std::vector< int > &axes, U init, OpS ops, OpC opc, Op op)
 
template<typename T , typename U , typename Op >
void mlx::core::reduction_op (const array &x, array &out, const std::vector< int > &axes, U init, Op op)