Compare commits

...

3 Commits

3 changed files with 24 additions and 21 deletions

View File

@ -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![],
}, },
})); }));

View File

@ -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`.

View File

@ -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)