MLX
 
Loading...
Searching...
No Matches
reduce.h File Reference

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 Shape &shape, const Strides &strides)
 
std::pair< Shape, Stridesmlx::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)