Use custom float-like class with numba

Numba Support: How do I do ...?
1

How can I let numba treat a custom class as a float? A minimal examle

import numba
import numpy as np

class MyFloat:    
    def __init__(self, f):
        self.f = f
        
    def __float__(self) -> float:       
        return float(f)
    
    # ... in the actual implementation more is defined to make the class work like a float

s = 1.1
x = np.array([1,2,3.])
m = MyFloat(1.12)

@numba.jit(nopython=True)
def f(x, a):
    return a + x**2


f(s, 2) # works!
f(x, 2) # works!
f(m, 2) # fails with TypingError: non-precise type pyobject

I can use forceobj=True to make the call f(m, 2) work, but that makes the calls to f(s, 2) and f(x, 2) very slow. Is there a way to specify that numba should treat objects of type MyFloat like a Python float (with conversion by calling __float__)?

I have read the Interval example at Example: an interval type — Numba 0.52.0.dev0+274.g626b40e-py3.7-linux-x86_64.egg documentation. But there it looks like one would have to teach numba about all the operations one could do with a MyFloat, instead of converting m to a float and then using the normal jitted method.

2

Just to clarify… the Interval example gives you the information you need?

3

@nelson2005 The Interval example does not provide the information I am looking for. I am looking for a way for Numba to treat my MyFloat class as a builtin float. So in this example:

@numba.jit(nopython=True)
def poly(x, a, b):
     return a * x + b

poly(1.1, 2.1, 3)
poly(1.1, MyFloat(2.1), 3)

The first call poly(1.1, 2.1, 3) is jitted by numba efficiencty. For the second call poly(1.1, MyFloat(2.1), 3) I would like numba to:

i) Convert MyFloat(2.1) to a python float (or a numba float64) using MyFloat.__float__
ii) Then call poly(1.1, 2.1, 3) the way numba normally handles two float arguments

Here both i) and ii) are fast operations, so I would hope numba can jit it much faster then the regular python call.

Note: writing poly(1.1, float(MyFloat(2.1)), 3) does not work, since I do not know beforehand the type of the argument a. It could also be a numpy array, in which case casting to float is not the right approach.