Is there anyway to wrap an imported numpy function into jax.numpy?

[PabloAMC]

I think this is a long shot, but: is there any way to convert a numpy function into a jax.numpy function? Just doing the usual trick of changing

import numpy as np

by

import jax.numpy as np

is not satisfactory because requires me to modify external libraries, and furthermore, sometimes I get errors such as, in lax_numpy.py,

@_wraps(np.array, lax_description=_ARRAY_DOC)
def array(object, dtype=None, copy=True, order="K", ndmin=0):
  if order is not None and order != "K":
    raise NotImplementedError("Only implemented for order='K'")

Thanks!

[jakevdp]

No, unfortunately there's no way to do this in general. JAX is unable to implement the entire numpy API... as a simple example, numpy will happily create and manipulate arrays of strings. However XLA (JAX's backend) has no string type, so it's not possible to express the corresponding operations in JAX.

You should think of jax.numpy as a numpy-like API, but not as a 100% faithful drop-in replacement for numpy.