lyken
/
nac3
forked from M-Labs/nac3
1
0
Fork 0
Code Pull Requests Activity

core: add np_argmax/np_argmin functions

This commit is contained in:
abdul124 2024-07-12 18:18:28 +08:00
parent d658d9b00e
commit cea7cade51
2 changed files with 32 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
nac3core/src/toplevel/builtins.rs
View File

@ -510,7 +510,10 @@ impl<'a> BuiltinBuilder<'a> {
PrimDef::FunMin | PrimDef::FunMax => self.build_min_max_function(prim), PrimDef::FunMin | PrimDef::FunMax => self.build_min_max_function(prim),
PrimDef::FunNpMin | PrimDef::FunNpMax => self.build_np_min_max_function(prim), PrimDef::FunNpArgmin
| PrimDef::FunNpArgmax
| PrimDef::FunNpMin
| PrimDef::FunNpMax => self.build_np_max_min_function(prim),
PrimDef::FunNpMinimum | PrimDef::FunNpMaximum => { PrimDef::FunNpMinimum | PrimDef::FunNpMaximum => {
self.build_np_minimum_maximum_function(prim) self.build_np_minimum_maximum_function(prim)
@ -1555,39 +1558,42 @@ impl<'a> BuiltinBuilder<'a> {
} }
} }
/// Build the functions `np_min()` and `np_max()`. /// Build the functions `np_max()`, `np_min()`, `np_argmax()` and `np_argmin()`
fn build_np_min_max_function(&mut self, prim: PrimDef) -> TopLevelDef { /// Calls `call_numpy_max_min` with the function name
debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpMin, PrimDef::FunNpMax]); fn build_np_max_min_function(&mut self, prim: PrimDef) -> TopLevelDef {
debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpArgmin, PrimDef::FunNpArgmax, PrimDef::FunNpMin, PrimDef::FunNpMax]);
let ret_ty = self.unifier.get_fresh_var(Some("R".into()), None); let (var_map, ret_ty) = match prim {
let var_map = self PrimDef::FunNpArgmax | PrimDef::FunNpArgmin => {
.num_or_ndarray_var_map (self.num_or_ndarray_var_map.clone(), self.primitives.int64)
.clone() },
.into_iter() PrimDef::FunNpMax | PrimDef::FunNpMin => {
.chain(once((ret_ty.id, ret_ty.ty))) let ret_ty = self.unifier.get_fresh_var(Some("R".into()), None);
.collect::<IndexMap<_, _>>(); let var_map = self
.num_or_ndarray_var_map
.clone()
.into_iter()
.chain(once((ret_ty.id, ret_ty.ty)))
.collect::<IndexMap<_, _>>();
(var_map, ret_ty.ty)
},
_ => unreachable!()
};
create_fn_by_codegen( create_fn_by_codegen(
self.unifier, self.unifier,
&var_map, &var_map,
prim.name(), prim.name(),
ret_ty.ty, ret_ty,
&[(self.float_or_ndarray_ty.ty, "a")], &[(self.num_or_ndarray_ty.ty, "a")],
Box::new(move |ctx, _, fun, args, generator| { Box::new(move |ctx, _, fun, args, generator| {
let a_ty = fun.0.args[0].ty; let a_ty = fun.0.args[0].ty;
let a = args[0].1.clone().to_basic_value_enum(ctx, generator, a_ty)?; let a = args[0].1.clone().to_basic_value_enum(ctx, generator, a_ty)?;
let func = match prim { Ok(Some(builtin_fns::call_numpy_max_min(generator, ctx, (a_ty, a), &prim.name())?))
PrimDef::FunNpMin => builtin_fns::call_numpy_min,
PrimDef::FunNpMax => builtin_fns::call_numpy_max,
_ => unreachable!(),
};
Ok(Some(func(generator, ctx, (a_ty, a))?))
}), }),
) )
} }
/// Build the functions `np_minimum()` and `np_maximum()`. /// Build the functions `np_minimum()` and `np_maximum()`.
fn build_np_minimum_maximum_function(&mut self, prim: PrimDef) -> TopLevelDef { fn build_np_minimum_maximum_function(&mut self, prim: PrimDef) -> TopLevelDef {
debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpMinimum, PrimDef::FunNpMaximum]); debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpMinimum, PrimDef::FunNpMaximum]);

4
nac3core/src/toplevel/helper.rs
View File

@ -62,9 +62,11 @@ pub enum PrimDef {
FunMin, FunMin,
FunNpMin, FunNpMin,
FunNpMinimum, FunNpMinimum,
FunNpArgmin,
FunMax, FunMax,
FunNpMax, FunNpMax,
FunNpMaximum, FunNpMaximum,
FunNpArgmax,
FunAbs, FunAbs,
FunNpIsNan, FunNpIsNan,
FunNpIsInf, FunNpIsInf,
@ -216,9 +218,11 @@ impl PrimDef {
PrimDef::FunMin => fun("min", None), PrimDef::FunMin => fun("min", None),
PrimDef::FunNpMin => fun("np_min", None), PrimDef::FunNpMin => fun("np_min", None),
PrimDef::FunNpMinimum => fun("np_minimum", None), PrimDef::FunNpMinimum => fun("np_minimum", None),
PrimDef::FunNpArgmin => fun("np_argmin", None),
PrimDef::FunMax => fun("max", None), PrimDef::FunMax => fun("max", None),
PrimDef::FunNpMax => fun("np_max", None), PrimDef::FunNpMax => fun("np_max", None),
PrimDef::FunNpMaximum => fun("np_maximum", None), PrimDef::FunNpMaximum => fun("np_maximum", None),
PrimDef::FunNpArgmax => fun("np_argmax", None),
PrimDef::FunAbs => fun("abs", None), PrimDef::FunAbs => fun("abs", None),
PrimDef::FunNpIsNan => fun("np_isnan", None), PrimDef::FunNpIsNan => fun("np_isnan", None),
PrimDef::FunNpIsInf => fun("np_isinf", None), PrimDef::FunNpIsInf => fun("np_isinf", None),