CompiledFunction switch to PjitFunction:

[rmoyard]

Hi Jax team!

I am encountering some issue with the new version 0.4.4. I have a class called DefaultQubitJax that inherits from an abstract class and we redefine some methods of the class like this:

_abs = staticmethod(jax.numpy.abs)
_conj = staticmethod(jax.numpy.conj)

Since 0.4.4 the device cannot be instantiated because of three different function abs, conj and dot, the error message is the following:

Can't instantiate abstract class DefaultQubitJax with abstract methods _abs, _conj, _dot

It makes me think that with the recent update, those three functions became abstract. I can tell that the base class changed if you do:

print(type(jax.numpy.abs), jax.numpy.abs)
print(type(jax.numpy.dot), jax.numpy.dot)

Before 0.4.4 you obtain:

<class 'jaxlib.xla_extension.CompiledFunction'> <CompiledFunction of <function absolute at 0x7f8d34bcd940>>
<class 'jaxlib.xla_extension.CompiledFunction'> <CompiledFunction of <function dot at 0x7f8d349c9820>>

After 0.4.4 you obtain:

<class 'jaxlib.xla_extension.PjitFunction'> <PjitFunction of <function absolute at 0x7f01f21dcaf0>>
<class 'jaxlib.xla_extension.PjitFunction'> <PjitFunction of <function dot at 0x7f01f1ff7280>>

What is the impact of the switch from CompiledFunction to PjitFunction, does it mean that the functions became abstract? Any work around to still use it to redefine abstract methods? It is probably some consequences of the jit and pjit merge.

[jakevdp]

Thanks for the question - this is very curious and I do not know why it's happening. For reference, here's a minimal repro that passes for the old jit (the default in jax 0.4.3 and older) and fails with the new pjit (the default in jax 0.4.4):

import abc
import jax

class Base(abc.ABC):
  @abc.abstractstaticmethod
  def _abs(x): ...

class Derived(Base):
  _abs = staticmethod(jax.numpy.abs)

d = Derived()

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-20-b8456e68a3b7>](https://localhost:8080/#) in <module>
      9   _abs = staticmethod(jax.numpy.abs)
     10 
---> 11 d = Derived()

TypeError: Can't instantiate abstract class Derived with abstract methods _abs

There must be some detail about how PjitFunction is implemented that causes Python to treat it as an abstract method, but I can't figure out what the difference might be. Perhaps @yashk2810 might know?

[jakevdp]

OK, so it looks like it has nothing to do with the pjit switch, but rather is due to #14474. It turns out that abstractmethod mutates its arguments in-place, and that change passes functions like jnp.abs to abstractmethod in order to dynamically populate jax.Array. I'll work on a fix.

[jakevdp]

Should be fixed by #14555

[rmoyard]

Thank you very much for the explanation and the quick fix 💯