Disclaimer: I already posted this on Reddit so apologies if it’s just reaching the same community

Hi Everyone, I having real trouble retaining the numpy array types under operations. I have defined the following type:
BatchNactionsNpFloatType = Annotated[
np.ndarray[tuple[int,int], np.float32], Literal["batch", "n_actions"]
]

I have two arrays defined with this type e.g.:
x:BatchNactionsNpFloatType = np.array([[1.0,2.0],[1.0,3.0],[1.0,2.0],[1.0,2.0]])
y:BatchNactionsNpFloatType = np.array([[1.0,2.0],[1.0,2.0],[3.0,2.0],[1.0,2.0]])
And I perform a simple operation:

res = np.equal(x,y)

However, according to VS code, ‘res’ is of type Any. I’m really confused why it wouldn’t return something like np.ndarray[np.bool_]?

Thanks!

I was expecting the np.equal operation to return an ndarray bool type