forked from M-Labs/nac3
core: categorize np_{transpose,reshape} as 'view functions'
This commit is contained in:
parent
a2a1e74be0
commit
5f94f4a4cd
|
@ -516,6 +516,10 @@ impl<'a> BuiltinBuilder<'a> {
|
||||||
self.build_ndarray_property_getter_function(prim)
|
self.build_ndarray_property_getter_function(prim)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
PrimDef::FunNpTranspose | PrimDef::FunNpReshape => {
|
||||||
|
self.build_ndarray_view_function(prim)
|
||||||
|
}
|
||||||
|
|
||||||
PrimDef::FunStr => self.build_str_function(),
|
PrimDef::FunStr => self.build_str_function(),
|
||||||
|
|
||||||
PrimDef::FunFloor | PrimDef::FunFloor64 | PrimDef::FunCeil | PrimDef::FunCeil64 => {
|
PrimDef::FunFloor | PrimDef::FunFloor64 | PrimDef::FunCeil | PrimDef::FunCeil64 => {
|
||||||
|
@ -581,10 +585,6 @@ impl<'a> BuiltinBuilder<'a> {
|
||||||
| PrimDef::FunNpHypot
|
| PrimDef::FunNpHypot
|
||||||
| PrimDef::FunNpNextAfter => self.build_np_2ary_function(prim),
|
| PrimDef::FunNpNextAfter => self.build_np_2ary_function(prim),
|
||||||
|
|
||||||
PrimDef::FunNpTranspose | PrimDef::FunNpReshape => {
|
|
||||||
self.build_np_sp_ndarray_function(prim)
|
|
||||||
}
|
|
||||||
|
|
||||||
PrimDef::FunNpDot
|
PrimDef::FunNpDot
|
||||||
| PrimDef::FunNpLinalgCholesky
|
| PrimDef::FunNpLinalgCholesky
|
||||||
| PrimDef::FunNpLinalgQr
|
| PrimDef::FunNpLinalgQr
|
||||||
|
@ -1440,6 +1440,57 @@ impl<'a> BuiltinBuilder<'a> {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Build np/sp functions that take as input `NDArray` only
|
||||||
|
fn build_ndarray_view_function(&mut self, prim: PrimDef) -> TopLevelDef {
|
||||||
|
debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpTranspose, PrimDef::FunNpReshape]);
|
||||||
|
|
||||||
|
match prim {
|
||||||
|
PrimDef::FunNpTranspose => {
|
||||||
|
let ndarray_ty = self.unifier.get_fresh_var_with_range(
|
||||||
|
&[self.ndarray_num_ty],
|
||||||
|
Some("T".into()),
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
create_fn_by_codegen(
|
||||||
|
self.unifier,
|
||||||
|
&into_var_map([ndarray_ty]),
|
||||||
|
prim.name(),
|
||||||
|
self.ndarray_num_ty,
|
||||||
|
&[(self.ndarray_num_ty, "x")],
|
||||||
|
Box::new(move |ctx, _, fun, args, generator| {
|
||||||
|
let arg_ty = fun.0.args[0].ty;
|
||||||
|
let arg_val =
|
||||||
|
args[0].1.clone().to_basic_value_enum(ctx, generator, arg_ty)?;
|
||||||
|
Ok(Some(ndarray_transpose(generator, ctx, (arg_ty, arg_val))?))
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NOTE: on `ndarray_factory_fn_shape_arg_tvar` and
|
||||||
|
// the `param_ty` for `create_fn_by_codegen`.
|
||||||
|
//
|
||||||
|
// Similar to `build_ndarray_from_shape_factory_function` we delegate the responsibility of typechecking
|
||||||
|
// to [`typecheck::type_inferencer::Inferencer::fold_numpy_function_call_shape_argument`],
|
||||||
|
// and use a dummy [`TypeVar`] `ndarray_factory_fn_shape_arg_tvar` as a placeholder for `param_ty`.
|
||||||
|
PrimDef::FunNpReshape => create_fn_by_codegen(
|
||||||
|
self.unifier,
|
||||||
|
&VarMap::new(),
|
||||||
|
prim.name(),
|
||||||
|
self.ndarray_num_ty,
|
||||||
|
&[(self.ndarray_num_ty, "x"), (self.ndarray_factory_fn_shape_arg_tvar.ty, "shape")],
|
||||||
|
Box::new(move |ctx, _, fun, args, generator| {
|
||||||
|
let x1_ty = fun.0.args[0].ty;
|
||||||
|
let x1_val = args[0].1.clone().to_basic_value_enum(ctx, generator, x1_ty)?;
|
||||||
|
let x2_ty = fun.0.args[1].ty;
|
||||||
|
let x2_val = args[1].1.clone().to_basic_value_enum(ctx, generator, x2_ty)?;
|
||||||
|
Ok(Some(ndarray_reshape(generator, ctx, (x1_ty, x1_val), (x2_ty, x2_val))?))
|
||||||
|
}),
|
||||||
|
),
|
||||||
|
|
||||||
|
_ => unreachable!(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Build the `str()` function.
|
/// Build the `str()` function.
|
||||||
fn build_str_function(&mut self) -> TopLevelDef {
|
fn build_str_function(&mut self) -> TopLevelDef {
|
||||||
let prim = PrimDef::FunStr;
|
let prim = PrimDef::FunStr;
|
||||||
|
@ -1927,57 +1978,6 @@ impl<'a> BuiltinBuilder<'a> {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build np/sp functions that take as input `NDArray` only
|
|
||||||
fn build_np_sp_ndarray_function(&mut self, prim: PrimDef) -> TopLevelDef {
|
|
||||||
debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpTranspose, PrimDef::FunNpReshape]);
|
|
||||||
|
|
||||||
match prim {
|
|
||||||
PrimDef::FunNpTranspose => {
|
|
||||||
let ndarray_ty = self.unifier.get_fresh_var_with_range(
|
|
||||||
&[self.ndarray_num_ty],
|
|
||||||
Some("T".into()),
|
|
||||||
None,
|
|
||||||
);
|
|
||||||
create_fn_by_codegen(
|
|
||||||
self.unifier,
|
|
||||||
&into_var_map([ndarray_ty]),
|
|
||||||
prim.name(),
|
|
||||||
self.ndarray_num_ty,
|
|
||||||
&[(self.ndarray_num_ty, "x")],
|
|
||||||
Box::new(move |ctx, _, fun, args, generator| {
|
|
||||||
let arg_ty = fun.0.args[0].ty;
|
|
||||||
let arg_val =
|
|
||||||
args[0].1.clone().to_basic_value_enum(ctx, generator, arg_ty)?;
|
|
||||||
Ok(Some(ndarray_transpose(generator, ctx, (arg_ty, arg_val))?))
|
|
||||||
}),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
// NOTE: on `ndarray_factory_fn_shape_arg_tvar` and
|
|
||||||
// the `param_ty` for `create_fn_by_codegen`.
|
|
||||||
//
|
|
||||||
// Similar to `build_ndarray_from_shape_factory_function` we delegate the responsibility of typechecking
|
|
||||||
// to [`typecheck::type_inferencer::Inferencer::fold_numpy_function_call_shape_argument`],
|
|
||||||
// and use a dummy [`TypeVar`] `ndarray_factory_fn_shape_arg_tvar` as a placeholder for `param_ty`.
|
|
||||||
PrimDef::FunNpReshape => create_fn_by_codegen(
|
|
||||||
self.unifier,
|
|
||||||
&VarMap::new(),
|
|
||||||
prim.name(),
|
|
||||||
self.ndarray_num_ty,
|
|
||||||
&[(self.ndarray_num_ty, "x"), (self.ndarray_factory_fn_shape_arg_tvar.ty, "shape")],
|
|
||||||
Box::new(move |ctx, _, fun, args, generator| {
|
|
||||||
let x1_ty = fun.0.args[0].ty;
|
|
||||||
let x1_val = args[0].1.clone().to_basic_value_enum(ctx, generator, x1_ty)?;
|
|
||||||
let x2_ty = fun.0.args[1].ty;
|
|
||||||
let x2_val = args[1].1.clone().to_basic_value_enum(ctx, generator, x2_ty)?;
|
|
||||||
Ok(Some(ndarray_reshape(generator, ctx, (x1_ty, x1_val), (x2_ty, x2_val))?))
|
|
||||||
}),
|
|
||||||
),
|
|
||||||
|
|
||||||
_ => unreachable!(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Build `np_linalg` and `sp_linalg` functions
|
/// Build `np_linalg` and `sp_linalg` functions
|
||||||
///
|
///
|
||||||
/// The input to these functions must be floating point `NDArray`
|
/// The input to these functions must be floating point `NDArray`
|
||||||
|
|
|
@ -56,6 +56,10 @@ pub enum PrimDef {
|
||||||
FunNpShape,
|
FunNpShape,
|
||||||
FunNpStrides,
|
FunNpStrides,
|
||||||
|
|
||||||
|
// NumPy ndarray view functions
|
||||||
|
FunNpTranspose,
|
||||||
|
FunNpReshape,
|
||||||
|
|
||||||
// Miscellaneous NumPy & SciPy functions
|
// Miscellaneous NumPy & SciPy functions
|
||||||
FunNpRound,
|
FunNpRound,
|
||||||
FunNpFloor,
|
FunNpFloor,
|
||||||
|
@ -103,8 +107,6 @@ pub enum PrimDef {
|
||||||
FunNpLdExp,
|
FunNpLdExp,
|
||||||
FunNpHypot,
|
FunNpHypot,
|
||||||
FunNpNextAfter,
|
FunNpNextAfter,
|
||||||
FunNpTranspose,
|
|
||||||
FunNpReshape,
|
|
||||||
|
|
||||||
// Linalg functions
|
// Linalg functions
|
||||||
FunNpDot,
|
FunNpDot,
|
||||||
|
@ -246,6 +248,10 @@ impl PrimDef {
|
||||||
PrimDef::FunNpShape => fun("np_shape", None),
|
PrimDef::FunNpShape => fun("np_shape", None),
|
||||||
PrimDef::FunNpStrides => fun("np_strides", None),
|
PrimDef::FunNpStrides => fun("np_strides", None),
|
||||||
|
|
||||||
|
// NumPy NDArray view functions
|
||||||
|
PrimDef::FunNpTranspose => fun("np_transpose", None),
|
||||||
|
PrimDef::FunNpReshape => fun("np_reshape", None),
|
||||||
|
|
||||||
// Miscellaneous NumPy & SciPy functions
|
// Miscellaneous NumPy & SciPy functions
|
||||||
PrimDef::FunNpRound => fun("np_round", None),
|
PrimDef::FunNpRound => fun("np_round", None),
|
||||||
PrimDef::FunNpFloor => fun("np_floor", None),
|
PrimDef::FunNpFloor => fun("np_floor", None),
|
||||||
|
@ -293,8 +299,6 @@ impl PrimDef {
|
||||||
PrimDef::FunNpLdExp => fun("np_ldexp", None),
|
PrimDef::FunNpLdExp => fun("np_ldexp", None),
|
||||||
PrimDef::FunNpHypot => fun("np_hypot", None),
|
PrimDef::FunNpHypot => fun("np_hypot", None),
|
||||||
PrimDef::FunNpNextAfter => fun("np_nextafter", None),
|
PrimDef::FunNpNextAfter => fun("np_nextafter", None),
|
||||||
PrimDef::FunNpTranspose => fun("np_transpose", None),
|
|
||||||
PrimDef::FunNpReshape => fun("np_reshape", None),
|
|
||||||
|
|
||||||
// Linalg functions
|
// Linalg functions
|
||||||
PrimDef::FunNpDot => fun("np_dot", None),
|
PrimDef::FunNpDot => fun("np_dot", None),
|
||||||
|
|
|
@ -179,6 +179,10 @@ def patch(module):
|
||||||
module.np_identity = np.identity
|
module.np_identity = np.identity
|
||||||
module.np_array = np.array
|
module.np_array = np.array
|
||||||
|
|
||||||
|
# NumPy NDArray view functions
|
||||||
|
module.np_transpose = np.transpose
|
||||||
|
module.np_reshape = np.reshape
|
||||||
|
|
||||||
# NumPy NDArray property getters
|
# NumPy NDArray property getters
|
||||||
module.np_shape = np.shape
|
module.np_shape = np.shape
|
||||||
module.np_strides = lambda ndarray: ndarray.strides
|
module.np_strides = lambda ndarray: ndarray.strides
|
||||||
|
@ -222,8 +226,6 @@ def patch(module):
|
||||||
module.np_ldexp = np.ldexp
|
module.np_ldexp = np.ldexp
|
||||||
module.np_hypot = np.hypot
|
module.np_hypot = np.hypot
|
||||||
module.np_nextafter = np.nextafter
|
module.np_nextafter = np.nextafter
|
||||||
module.np_transpose = np.transpose
|
|
||||||
module.np_reshape = np.reshape
|
|
||||||
|
|
||||||
# SciPy Math functions
|
# SciPy Math functions
|
||||||
module.sp_spec_erf = special.erf
|
module.sp_spec_erf = special.erf
|
||||||
|
|
Loading…
Reference in New Issue