stan-dev / stanc3

The Stan transpiler (from Stan to C++ and beyond).
BSD 3-Clause "New" or "Revised" License
138 stars 44 forks source link

Typechecker: Allow predicate-based typechecking of library functions #1373

Open WardBrian opened 8 months ago

WardBrian commented 8 months ago

This would allow Array functions to be able to accept arrays of tuples, e.g. size

And also allow vectorization of things like a 7-parameter function, which would currently be too expensive to enumerate

WardBrian commented 8 months ago

Right now an argument is specified as UnsizedType.autodifftype * UnsizedType.t

I think we could define another type which looks like:

type argtype = 
  | Concrete of UnsizedType.t
  | Predicate of (UnsizedType.t -> bool) (* think of this like a C++ concept *)

We'd need to rewrite a fair amount of stuff to handle this but I don't think it would be too bad. The argument type of size would then be something like:

Predicate (function UArray _ -> true | _ -> false)

Unfortunately this becomes very difficult to print. We could have Predicate hold both UnsizedType.t -> bool and a string which we use as it's pretty-printed name, so like Predicate ((function UArray _ -> true | _ -> false), "array[] T")

SteveBronder commented 6 months ago

I think we could do this with an API like

(*Function that depends on two types*)
let matching_array_types x y = 
match x, y with 
  | UnsizedType.UArray _, UnsizedType.UArray _ when x == y -> Valid 
  | _ -> Invalid
Fun ("append_array",  matching_array_types)

(*Anything is valid*)
let all_valid _ = Valid
Fun ("size",  all_valid)

(*For the distributions*)
let is_vectorizable x = 
  match x with 
  | UnsizedType.UArray (UReal, UInt) | UVector | URowVector | UReal | UInt -> true
  | _ -> false 
in
let is_all_vectorizable x, y, z = 
  if is_vectorizable x && is_vectorizable y && is_vectorizable z the
    Valid
  else 
    Invalid
in
Fun ("normal_lpdf", is_all_vectorizable)
WardBrian commented 5 months ago

Some further thoughts: