numpy backend requires jax?

[hmaarrfk]

It seems that the numpy backend imports jax

keras/keras/backend/numpy/image.py

Line 1 in 9c675a9

import jax

and seems to use it below:

    return np.array(
        jax.image.resize(image, size, method=interpolation, antialias=antialias)
    )

is that intention?

[SuryanarayanaY]

It seems numpy backend has Jax dependency. I can find it in backend/numpy/nn.py also.

[fchollet]

Yes, that's intentional. We found that when implementing nn ops like conv or pooling, a low-level pure numpy implementation was vastly less performant than jax, so we just use jax for these ops.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[hmaarrfk]

I feel like this is seriously misleading. numpy and Jax offer drastically different compilation, testing, and integration challenges.

installing Jax and installing numpy can be quite different in resource limited environments.

I think that:

1. users of the pure numpy backend should expect slower, albeit numerically equivalent/similar to that of other more accelerated backends.
2. if you want, I think a hybrid approach if "trying to use Jax, but falling back to the much slower numpy operators" may be acceptable, but it may create surprising behavior.

for example, I would expect a backend of

1. "numpy" not to depend on Jax at all.
2. "numpy+Jax" to use Jax in select locations
3. "jax" primarily use Jax

as is, users cannot use the numpy backend at all without Jax. in my mind, it just shouldn't be called "numpy backend"

I am not saying this in any hostile way, just as a packager that tried to test things with numpy only and was quite surprised that I needed a heavy dependency like Jax to get any operators to work.

If my tone is negative, it is because in my mind, you hastily closed this issue without giving the chance to respond.
would it be acceptable to reopen this issue and at least acknowledge to others that this is a present limitation of the 3.0.0 version? Closing this issue just seems wrong.

[mthiboust]

I second the comment of @hmaarrfk. I think it would be useful to have the possibility of having a lightweight keras installation powered by the numpy backend in order to run small inference tasks in limited environments.

[fchollet]

That's a reasonable view. There's definitely a huge slowdown however.

If you'd like to open a PR to revert some ops to pure numpy implementation, we can consider it.

[mthiboust]

Thanks François for your answer. I see that the functions involved are jax.image.resize, jax.numpy.fft.* and jax.lax.conv_*. Don't know yet if I'll work on this PR.

[fchollet]

The fft ops should be a trivial change since they're already available in numpy.