Cannot patch __call__ on an object with vectorcall defined

[davidsaxton]

It appears that when the type of an object has JAX_TPFLAGS_HAVE_VECTORCALL set on tp_flags, then overrides to call from within python are ignored.

• PyObject_Call checks if vectorcall is available (which is only a property of the type), and if so, then does the vectorcall: https://github.com/python/cpython/blob/3.11/Objects/call.c#L328
• The implementation of vectorcall does not check for tp_call (is this where call ends up?) https://github.com/python/cpython/blob/3.11/Include/internal/pycore_call.h#L39

I encountered this bug when using the latest version of JAX, which introduced vectorcalls in tensorflow/tensorflow@bf3eb11 :

import jax
g = jax.jit(lambda x: x + 1)
print(type(g))
# > <class 'google3.third_party.tensorflow.compiler.xla.python.xla_extension.CompiledFunction'>
print(g(5))
# > 6
type(g).__call__ = lambda *args: 0
print(g(5))
# > 6

[encukou]

The class is responsible for ensuring that that vectorcall matches tp_call. Unfortunately, that's hard to do for mutable classes, where tp_call can be overridden by users. Consider making the class immutable.
It'll be changed in 3.12, see the change note: https://docs.python.org/3.12/c-api/call.html#the-vectorcall-protocol

[encukou]

Also see this warning: https://docs.python.org/3/c-api/typeobj.html?highlight=vectorcall_offset#c.PyTypeObject.tp_vectorcall_offset
PRs to make the docs clearer would be welcome.