Recommended approach for binding functions which have an argument with multiple types

[awni]

Suppose you want to have a function which takes an argument x which could be some user defined type (array in our case) or just a plain int.

There are several ways of doing this in nanobind that I know of:

1. Use overloads
2. Use a variant std::variant<int, array>
3. Get a nb::object and try casting it

I'm curious, are there any performance trade-offs associated with these? Any other things to think about when deciding between
them?

Here's a more detailed example. Using a variant we could bind the function like this:

m.def(                                                                                                                                                       
    "myfunction",  
    [](const std::variant<int, array>& x) {
      if (auto pv = std::get_if<int>(&x); pv) {                                                                                                              
        return function(*pv);
      } else {                                                                                                                                               
        return function(std::get<array>(x));                                                                                                
      }
    },  
...

Using nb::object we could do this:

m.def(                                                                                                                                                       
    "myfunction",  
    [](const nb::object& x) {                                                                                                         
        if (nb::isinstance<int>(x)) {                                                                                                                     
          return function(nb::cast<int>(x));                                                                                                        
        } else if (nb::isinstance<array>(x)) {                                                                                                    
          return function(nb::cast<array>(x));                                                                                              
        } else {                                                                                                                                             
          throw std::invalid_argument(                                                                                                                       
              "[function] Received invalid type for second input.");                                                                                       
        }                                                                                                                                                    
    },
...

We usually use the 2nd or 3rd option in MLX as it let's us control the function signatures and documentation. I also like the 3rd option because we can control the error message (bad_cast is confusing to people). But just wondering what other trade-offs there might be that we don't know about.

[wjakob]

They are all good 🤗. It mainly depends on your preference with respect to function signatures in tooling.

If you have a function with many overloads and many arguments (or few, but ones that are costly to cast, e.g. list -> std::vector<int> with big lists), then its likely cheaper to crunch it down to a single overload and work out the specific way to handle each argument in C++. Otherwise, nanobind will try the overloads one by one, which may involve some redundant type casting overheads.

[awni]

Thanks, that's been my working assumption, glad to have it confirmed!