Compare commits
3 Commits
abbaa506ad
...
3e92c491f5
Author | SHA1 | Date |
---|---|---|
David Mak | 3e92c491f5 | |
lyken | 7f629f1579 | |
lyken | 5640a793e2 |
|
@ -1550,36 +1550,29 @@ impl<'a> Inferencer<'a> {
|
||||||
}
|
}
|
||||||
// 2-argument ndarray n-dimensional creation functions
|
// 2-argument ndarray n-dimensional creation functions
|
||||||
if id == &"np_full".into() && args.len() == 2 {
|
if id == &"np_full".into() && args.len() == 2 {
|
||||||
let ExprKind::List { elts, .. } = &args[0].node else {
|
// Parse arguments
|
||||||
return report_error(
|
let shape_expr = args.remove(0);
|
||||||
format!(
|
let (ndims, shape) =
|
||||||
"Expected List literal for first argument of {id}, got {}",
|
self.fold_numpy_function_call_shape_argument(*id, 0, shape_expr)?; // Special handling for `shape`
|
||||||
args[0].node.name()
|
|
||||||
)
|
|
||||||
.as_str(),
|
|
||||||
args[0].location,
|
|
||||||
);
|
|
||||||
};
|
|
||||||
|
|
||||||
let ndims = elts.len() as u64;
|
let fill_value = self.fold_expr(args.remove(0))?;
|
||||||
|
|
||||||
let arg0 = self.fold_expr(args.remove(0))?;
|
// Build the return type
|
||||||
let arg1 = self.fold_expr(args.remove(0))?;
|
let dtype = fill_value.custom.unwrap();
|
||||||
|
|
||||||
let ty = arg1.custom.unwrap();
|
|
||||||
let ndims = self.unifier.get_fresh_literal(vec![SymbolValue::U64(ndims)], None);
|
let ndims = self.unifier.get_fresh_literal(vec![SymbolValue::U64(ndims)], None);
|
||||||
let ret = make_ndarray_ty(self.unifier, self.primitives, Some(ty), Some(ndims));
|
let ret = make_ndarray_ty(self.unifier, self.primitives, Some(dtype), Some(ndims));
|
||||||
|
|
||||||
let custom = self.unifier.add_ty(TypeEnum::TFunc(FunSignature {
|
let custom = self.unifier.add_ty(TypeEnum::TFunc(FunSignature {
|
||||||
args: vec![
|
args: vec![
|
||||||
FuncArg {
|
FuncArg {
|
||||||
name: "shape".into(),
|
name: "shape".into(),
|
||||||
ty: arg0.custom.unwrap(),
|
ty: shape.custom.unwrap(),
|
||||||
default_value: None,
|
default_value: None,
|
||||||
is_vararg: false,
|
is_vararg: false,
|
||||||
},
|
},
|
||||||
FuncArg {
|
FuncArg {
|
||||||
name: "fill_value".into(),
|
name: "fill_value".into(),
|
||||||
ty: arg1.custom.unwrap(),
|
ty: fill_value.custom.unwrap(),
|
||||||
default_value: None,
|
default_value: None,
|
||||||
is_vararg: false,
|
is_vararg: false,
|
||||||
},
|
},
|
||||||
|
@ -1597,7 +1590,7 @@ impl<'a> Inferencer<'a> {
|
||||||
location: func.location,
|
location: func.location,
|
||||||
node: ExprKind::Name { id: *id, ctx: *ctx },
|
node: ExprKind::Name { id: *id, ctx: *ctx },
|
||||||
}),
|
}),
|
||||||
args: vec![arg0, arg1],
|
args: vec![shape, fill_value],
|
||||||
keywords: vec![],
|
keywords: vec![],
|
||||||
},
|
},
|
||||||
}));
|
}));
|
||||||
|
|
|
@ -670,8 +670,8 @@ impl Unifier {
|
||||||
let num_args = posargs.len() + kwargs.len();
|
let num_args = posargs.len() + kwargs.len();
|
||||||
|
|
||||||
// Now we check the arguments against the parameters,
|
// Now we check the arguments against the parameters,
|
||||||
// and depending on what `call_info` is, we might change how the behavior `unify_call()`
|
// and depending on what `call_info` is, we might change how `unify_call()` behaves
|
||||||
// in hopes to improve user error messages when type checking fails.
|
// to improve user error messages when type checking fails.
|
||||||
match operator_info {
|
match operator_info {
|
||||||
Some(OperatorInfo::IsBinaryOp { self_type, operator }) => {
|
Some(OperatorInfo::IsBinaryOp { self_type, operator }) => {
|
||||||
// The call is written in the form of (say) `a + b`.
|
// The call is written in the form of (say) `a + b`.
|
||||||
|
|
|
@ -114,12 +114,22 @@ def test_ndarray_ones():
|
||||||
n: ndarray[float, 1] = np_ones([1])
|
n: ndarray[float, 1] = np_ones([1])
|
||||||
output_ndarray_float_1(n)
|
output_ndarray_float_1(n)
|
||||||
|
|
||||||
|
dim = (1,)
|
||||||
|
n_tup: ndarray[float, 1] = np_ones(dim)
|
||||||
|
output_ndarray_float_1(n_tup)
|
||||||
|
|
||||||
def test_ndarray_full():
|
def test_ndarray_full():
|
||||||
n_float: ndarray[float, 1] = np_full([1], 2.0)
|
n_float: ndarray[float, 1] = np_full([1], 2.0)
|
||||||
output_ndarray_float_1(n_float)
|
output_ndarray_float_1(n_float)
|
||||||
n_i32: ndarray[int32, 1] = np_full([1], 2)
|
n_i32: ndarray[int32, 1] = np_full([1], 2)
|
||||||
output_ndarray_int32_1(n_i32)
|
output_ndarray_int32_1(n_i32)
|
||||||
|
|
||||||
|
dim = (1,)
|
||||||
|
n_float_tup: ndarray[float, 1] = np_full(dim, 2.0)
|
||||||
|
output_ndarray_float_1(n_float_tup)
|
||||||
|
n_i32_tup: ndarray[int32, 1] = np_full(dim, 2)
|
||||||
|
output_ndarray_int32_1(n_i32_tup)
|
||||||
|
|
||||||
def test_ndarray_eye():
|
def test_ndarray_eye():
|
||||||
n: ndarray[float, 2] = np_eye(2)
|
n: ndarray[float, 2] = np_eye(2)
|
||||||
output_ndarray_float_2(n)
|
output_ndarray_float_2(n)
|
||||||
|
|
Loading…
Reference in New Issue