Jax backend by eickenberg · Pull Request #757 · kymatio/kymatio

eickenberg · 2021-06-30T01:07:58Z

Here is the jax backend for 1D. Tests pass.

I will add 2D and 3D here in the next days.

Please start reviewing this when you get a chance.

janden

Great work! Some things do discuss here. Mainly clarifying the triple inheritance stuff which scare me a bit.

janden · 2021-06-30T11:59:34Z

+from .numpy_frontend import ScatteringNumPy1D
+from .base_frontend import ScatteringBase1D
+
+class ScatteringJax1D(ScatteringJax, ScatteringNumPy1D, ScatteringBase1D):


Again, slightly terrifying and definitely deserving of a comment if not a rewrite.

Explanation:

Multiple inheritance generally looks for methods/attributes in the order of the classes specified above. Hence ScatteringJax goes first in order to have priority adding its attributes in.

Apart from the inheritance of ScatteringNumpPy1D this class and its __init__ look exactly like the one in ScatteringNumPy1D itself.

I had to include ScatteringBase1D because I needed to overwrite the __init__, basically copying and pasting the __init__ from ScatteringNumPy1D but with the backend changed. I had to do that because otherwise it was doing something numpy-specific from the ScatteringNumPy1D.__init__. In order to do that overwriting I needed access to ScatteringBase1D (from which ScatteringNumPy1D also inherits and which it uses in its __init__)

I initially tried switching positions 2 and 3 and I got a multiple inheritance error (because ScatteringNumPy1D derives from ScatteringBase1D) and so not all permutations are actually allowed.

Will add some comments

added some comments

For 3. I don't know that you need to include ScatteringBase1D in the inheritance list just to call its constructor. It is a super-super-class of the current one, so there's no issue calling its constructor like you're doing. I tried removing this locally and tests pass.

ScatteringBase1D is already inherited via ScatteringNumpy1D. Its redundant here.

@MuawizChaudhary update inheritance?

janden · 2021-06-30T12:02:51Z

+    assert np.allclose(jax_result, numpy_result)
+
+
+def test_jnp_multiplication():


How is this different from the previous test?

don't know! I copied and pasted from the previous PR :D Will check

janden · 2021-06-30T12:04:46Z

@@ -0,0 +1,36 @@
+import pytest


Rename test_jax_scattering1d.py per convention.

janden · 2021-06-30T12:06:19Z

+    T = x.shape[-1]
+    scattering = Scattering1D(J, T, Q, backend=backend, frontend='jax')
+
+    Sx = scattering(x) 


Do we know that Jax is actually being used here (as opposed to NP)? There should be a way to find out using the mock package, I think. Conversely, it would be good to make sure that NP is not being used (except for filter generation) so that a call doesn't slip through by accident.

Here is another pictorial proof.

You mean build a fake numpy or something that errors on all/most methods? With https://docs.python.org/dev/library/unittest.mock.html ?

You mean build a fake numpy or something that errors on all/most methods?

Yeah or is that overkill, do you think? I don't know that you need to replace numpy (you could do that with monkeypatch, I think). You can just do things like .called on various methods. Maybe NP will crash already if you give it Jax arrays.

eickenberg

Thanks for the review!

I left some comments about your questions.

Going to address them code-side now

eickenberg · 2021-07-01T01:39:49Z

+from .numpy_frontend import ScatteringNumPy1D
+from .base_frontend import ScatteringBase1D
+
+class ScatteringJax1D(ScatteringJax, ScatteringNumPy1D, ScatteringBase1D):


Explanation:

Multiple inheritance generally looks for methods/attributes in the order of the classes specified above. Hence ScatteringJax goes first in order to have priority adding its attributes in.

Apart from the inheritance of ScatteringNumpPy1D this class and its __init__ look exactly like the one in ScatteringNumPy1D itself.

I had to include ScatteringBase1D because I needed to overwrite the __init__, basically copying and pasting the __init__ from ScatteringNumPy1D but with the backend changed. I had to do that because otherwise it was doing something numpy-specific from the ScatteringNumPy1D.__init__. In order to do that overwriting I needed access to ScatteringBase1D (from which ScatteringNumPy1D also inherits and which it uses in its __init__)

I initially tried switching positions 2 and 3 and I got a multiple inheritance error (because ScatteringNumPy1D derives from ScatteringBase1D) and so not all permutations are actually allowed.

Will add some comments

eickenberg · 2021-07-01T01:43:19Z

+    assert np.allclose(jax_result, numpy_result)
+
+
+def test_jnp_multiplication():


don't know! I copied and pasted from the previous PR :D Will check

eickenberg · 2021-07-01T01:47:05Z

@@ -0,0 +1,36 @@
+import pytest


eickenberg · 2021-07-01T01:48:49Z

+    T = x.shape[-1]
+    scattering = Scattering1D(J, T, Q, backend=backend, frontend='jax')
+
+    Sx = scattering(x) 


Here is another pictorial proof.

You mean build a fake numpy or something that errors on all/most methods? With https://docs.python.org/dev/library/unittest.mock.html ?

eickenberg · 2021-07-01T01:49:41Z

+    test_data_dir = os.path.dirname(__file__)
+
+    with open(os.path.join(test_data_dir, 'test_data_1d.npz'), 'rb') as f:
+        buffer = io.BytesIO(f.read())


huh, why were we doing this? Instead of np.load?

eickenberg · 2021-07-01T19:57:21Z

Added 2D jax backend with passing tests. This required moving some of the numpy input checks from the scattering function to the numpy backend

eickenberg · 2021-07-01T21:30:56Z

Added 3D backend.

Another quick review and discussion of the inheritance situation would be great!

janden · 2021-07-02T08:41:20Z

Ok I'll review 2D and 3D once we've converged on 1D.

eickenberg · 2022-02-14T23:58:49Z

@janden could we touch base on this one soon-ish and try to converge on a design decision?

MuawizChaudhary · 2022-05-01T19:52:39Z

Whats the current state of this PR @janden @eickenberg ? I can jump in and help. Sorry I never finished this during my internship!

lostanlen · 2022-05-02T14:23:14Z

I don't have the relevant background in JAX but i heartily support this: JAX looks really easy to use as a drop-in replacement for NumPy !

MuawizChaudhary

tests error out on the numpy 3d. Why did the input checks have to be changed for numpy again?

MuawizChaudhary · 2022-05-31T17:15:12Z

+from .numpy_frontend import ScatteringNumPy1D
+from .base_frontend import ScatteringBase1D
+
+class ScatteringJax1D(ScatteringJax, ScatteringNumPy1D, ScatteringBase1D):


ScatteringBase1D is already inherited via ScatteringNumpy1D. Its redundant here.

MuawizChaudhary · 2022-06-03T01:00:45Z

-        self.backend.input_checks(x)
-


Why this removal?

i'll look into it later

MuawizChaudhary

I can fix some of this up for y'all if thats fine with yall.

MuawizChaudhary · 2022-06-16T05:49:06Z

@janden I fixed up various things in the jax front/backends. So this PR is ready for your review. Unfortunately, it appears that Jenkins does not have Jax installed. So that needs fixing.

MuawizChaudhary

I think we could merge this in at any time. Unsure if we wanna do it soon tho, ot later closer to release.

lostanlen · 2022-06-18T22:41:40Z

DEPRECATION: Python 3.5 reached the end of its life on September 13th, 2020. Please upgrade your Python as Python 3.5 is no longer maintained. pip 21.0 will drop support for Python 3.5 in January 2021. pip 21.0 will remove support for this functionality.

shall we remove py35 from Pip CI ?

MuawizChaudhary · 2022-06-19T04:05:25Z

DEPRECATION: Python 3.5 reached the end of its life on September 13th, 2020. Please upgrade your Python as Python 3.5 is no longer maintained. pip 21.0 will drop support for Python 3.5 in January 2021. pip 21.0 will remove support for this functionality.

shall we remove py35 from Pip CI ?

I vote yes, can you make an issue.

lostanlen · 2022-06-19T08:06:11Z

I have attempted to update the GitHub worfklow in a recent PR (#892 )
I would say that JAX is next in line after this

janden · 2022-06-21T08:05:54Z

I agree. We'll rebase this once #892 is merged, then get back to reviewing it.

MuawizChaudhary · 2022-06-21T13:21:57Z

@janden i think you have the ability to reset tests for a pr?

janden · 2022-06-21T13:23:27Z

@janden i think you have the ability to reset tests for a pr?

Yes but it won't make a difference until we've rebased.

lostanlen · 2022-06-21T14:27:28Z

Alright, i'm rebasing. Everyone stand back

lostanlen · 2022-06-21T14:37:21Z

Jenkins is failing ... Torch 2D ?

tests/scattering2d/test_torch_scattering2d.py ...............sSending interrupt signal to process
Terminated
script returned exit code 143

janden · 2022-06-21T18:32:11Z

Jenkins is failing ... Torch 2D ?

tests/scattering2d/test_torch_scattering2d.py ...............sSending interrupt signal to process
Terminated
script returned exit code 143

My guess is that rebuilding the docker image took too much time and there wasn't enough time to run the tests. I could be wrong though. Re-running to check.

janden · 2022-06-21T18:39:12Z

Something happened with the rebase. There are lots of old commits that shouldn't be here. Do you have access to the original tree?

janden · 2022-06-21T18:45:06Z

Ok looks like I was able to restore it (see #898). Is it ok if I force-push that branch on here so that we can continue the review in the same place?

janden · 2022-06-21T19:24:41Z

Ok looks like I was able to restore it (see #898). Is it ok if I force-push that branch on here so that we can continue the review in the same place?

Ping @lostanlen

lostanlen · 2022-06-21T19:44:25Z

@janden yes, it's all good to me. apologies for the mishap in rebase

janden

Looks good. Mostly comments and questions about testing here, which I think can be streamlined. Since the implementation so closely mirrors the NP frontend, I think we can get away with a simple test against sample data for each dimension and be done (perhaps checking on host and on device to be safe). This will ensure that Jax is doing the correct thing. All the other tests basically reproduce the tests we have for the NP frontend.

janden · 2022-06-21T19:54:04Z

+        buffer = io.BytesIO(f.read())
+        data = np.load(buffer)
+
+    x = data['x']


Shouldn't we be casting to JNP arrays or doing device_put or something here (like in 1D)?

we should also be parameterizing Jax with cpu and gpu device in that instance then...

codecov-commenter · 2022-06-21T20:11:38Z

Codecov Report

Merging #757 (fa1a0df) into dev (6fc8547) will increase coverage by 0.20%.
The diff coverage is 90.75%.

@@            Coverage Diff             @@
##              dev     #757      +/-   ##
==========================================
+ Coverage   87.86%   88.07%   +0.20%     
==========================================
  Files          64       73       +9     
  Lines        2250     2297      +47     
==========================================
+ Hits         1977     2023      +46     
- Misses        273      274       +1

Flag	Coverage Δ
jenkins_main	`88.07% <90.75%> (+0.20%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
kymatio/frontend/numpy_frontend.py	`93.33% <ø> (-0.42%)`	⬇️
kymatio/jax.py	`0.00% <0.00%> (ø)`
kymatio/backend/numpy_backend.py	`97.56% <50.00%> (-2.44%)`	⬇️
kymatio/scattering2d/frontend/base_frontend.py	`92.15% <69.23%> (-7.85%)`	⬇️
kymatio/scattering1d/frontend/base_frontend.py	`93.26% <83.33%> (-4.17%)`	⬇️
kymatio/backend/jax_backend.py	`90.00% <90.00%> (ø)`
kymatio/frontend/base_frontend.py	`84.00% <100.00%> (+6.22%)`	⬆️
kymatio/frontend/entry.py	`100.00% <100.00%> (ø)`
kymatio/frontend/jax_frontend.py	`100.00% <100.00%> (ø)`
kymatio/scattering1d/backend/jax_backend.py	`100.00% <100.00%> (ø)`
... and 18 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update c2a627f...fa1a0df. Read the comment docs.

lostanlen · 2022-06-21T21:29:32Z

@MuawizChaudhary we're almost there. would you like to help us with simplifying the tests here?

MuawizChaudhary · 2022-06-21T21:40:14Z

yeah, I'm at CVPR rn but I'll get to it tonight

MuawizChaudhary · 2022-06-22T03:58:37Z

... i fucked up

janden · 2022-06-22T05:17:21Z

Cherry-picked your commits onto the cleaned-up branch. Just remember to rebase your branch before you commit and push next time.

eickenberg

Looks great to me.

I had some questions wherever numpy files were modified, wondering whether that was from some previous PR or whether it pertained to making the structure jax compatible. I don't think any changes need to be made, but I'd like to understand the instances where the numpy files are modified in terms of this question, so comments appreciated.

I would also add an "Approve" to this PR, but it's grayed out for me, because I am the original author of this PR 😹

eickenberg · 2022-06-23T13:15:19Z

    def __call__(self, x):
        """This method is an alias for `scattering`."""

-        self.backend.input_checks(x)


Is this also from another PR or added during this one?

also from here

eickenberg · 2022-06-23T13:19:19Z

        ScatteringBase2D.create_filters(self)

    def scattering(self, input):
-        if not type(input) is np.ndarray:


again the question whether this numpy file modification is an integral part of this PR

eickenberg · 2022-06-23T13:22:14Z

+        integrals = []
        for i_q, q in enumerate(integral_powers):
-            integrals[:, i_q] = (input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1)
+            integrals.append((input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1))


is this modification to accommodate jax, because jax doesnt allow setting subarrays?
this will likely double the memory usage - worth considering whether to override this only in jax and keep numpy as it was?

yes, we have to do some .set and one additional operation.

eickenberg · 2022-06-23T13:22:37Z

        for i_q, q in enumerate(integral_powers):
-            integrals[:, i_q] = (input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1)
+            integrals.append((input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1))
+        integrals = cls._np.float32(cls._np.stack(integrals, axis=-1))


why is the cast to float32 necessary? isn't the input also float32?

it could be float64, so tests fail in that instance when we expect float32. Also, previously in this code integrals was defined as a float32 array.

eickenberg · 2022-06-23T13:26:53Z

    def scattering(self, input_array):
-        if not type(input_array) is np.ndarray:
-            raise TypeError('The input should be a NumPy array.')
+        self.backend.input_checks(input_array)


this is to let jax arrays through as well?
I like the change, just need to understand the motivation. I would have assumed that since in the jax case, np = jax.numpy, it should all be transparent

Yes. I would have assumed so too. But I think np.array should be valid. perhaps we should add a special check for numpy or jax in the jax backend, overriding input_checks?

MuawizChaudhary · 2022-06-23T15:16:17Z

    def __call__(self, x):
        """This method is an alias for `scattering`."""

-        self.backend.input_checks(x)


also from here

MuawizChaudhary · 2022-06-23T15:16:35Z

        ScatteringBase2D.create_filters(self)

    def scattering(self, input):
-        if not type(input) is np.ndarray:


MuawizChaudhary · 2022-06-23T15:18:00Z

+    def input_checks(x):
+        if x is None:
+            raise TypeError('The input should be not empty.')


@eickenberg you added this here.

perhaps we should check that its a jax array?

MuawizChaudhary · 2022-06-23T15:21:01Z

        for i_q, q in enumerate(integral_powers):
-            integrals[:, i_q] = (input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1)
+            integrals.append((input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1))
+        integrals = cls._np.float32(cls._np.stack(integrals, axis=-1))


it could be float64, so tests fail in that instance when we expect float32. Also, previously in this code integrals was defined as a float32 array.

MuawizChaudhary · 2022-06-23T15:21:23Z

+        integrals = []
        for i_q, q in enumerate(integral_powers):
-            integrals[:, i_q] = (input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1)
+            integrals.append((input_array ** q).reshape((input_array.shape[0], -1)).sum(axis=1))


yes, we have to do some .set and one additional operation.

MuawizChaudhary · 2022-06-23T15:23:13Z

    def scattering(self, input_array):
-        if not type(input_array) is np.ndarray:
-            raise TypeError('The input should be a NumPy array.')
+        self.backend.input_checks(input_array)


Yes. I would have assumed so too. But I think np.array should be valid. perhaps we should add a special check for numpy or jax in the jax backend, overriding input_checks?

janden

Great work!

…-like ndarrays

* ENH jax base backend and frontend * ENH jax backend and frontend for 1D * TST general jax.numpy - numpy correspondence tests * TST test for jax 1d backend * WIP address Joakim's comments * WIP jax 2d * WIP typo in jax 2d * ENH add 2D jax backend * ENH jax backend 3d * ENH jax entry point * MAINT fix 3d numpy frontend test * MAINT remove uneeded inheritances * STY change script names to be more inline with current script naming convention * TST edited workflows to account for jax and jaxlib * TST adding jax to jenkins * MAINT fix jaxbackend not complying with numpy backend after mergeal * MAINT cast to float32 * update JAX tests (tuple Q) * TST Fix seed for 2D FFT test (Jax) * MAINT removed lonely jax import * TST just have one test * TST removed unneeded test * MAINT make sure that we check for numpy ndarrays in addition to numpy-like ndarrays Co-authored-by: Muawiz Chaudhary <muawizc@gmail.com> Co-authored-by: Vincent Lostanlen <vincent.lostanlen@ls2n.fr> Co-authored-by: Joakim Andén <janden@kth.se>

eickenberg requested review from MuawizChaudhary, edouardoyallon and janden June 30, 2021 01:07

janden requested changes Jun 30, 2021

View reviewed changes

eickenberg commented Jul 1, 2021

View reviewed changes

Comment thread tests/scattering1d/test_jax_scattering1d.py

MuawizChaudhary reviewed Jun 3, 2022

View reviewed changes

MuawizChaudhary requested changes Jun 3, 2022

View reviewed changes

MuawizChaudhary requested a review from janden June 16, 2022 05:35

MuawizChaudhary changed the title ~~[WIP] Jax backend~~ Jax backend Jun 17, 2022

MuawizChaudhary approved these changes Jun 17, 2022

View reviewed changes

MuawizChaudhary self-assigned this Jun 17, 2022

lostanlen requested a review from MuawizChaudhary June 21, 2022 14:37

TST edited workflows to account for jax and jaxlib

3a3bfc3

MuawizChaudhary and others added 5 commits June 21, 2022 20:49

TST adding jax to jenkins

312a0b5

MAINT fix jaxbackend not complying with numpy backend after mergeal

fe7b06c

MAINT cast to float32

2957dc5

update JAX tests (tuple Q)

3f3d6a0

TST Fix seed for 2D FFT test (Jax)

686d04b

janden force-pushed the jax_backend branch from d5f1f2a to 686d04b Compare June 21, 2022 19:45

janden requested changes Jun 21, 2022

View reviewed changes

MuawizChaudhary force-pushed the jax_backend branch from 686d04b to ad4dc25 Compare June 22, 2022 03:57

MuawizChaudhary added 3 commits June 22, 2022 07:16

MAINT removed lonely jax import

1f83a49

TST just have one test

5277c3f

TST removed unneeded test

7d3bcf5

janden force-pushed the jax_backend branch from fa1a0df to 7d3bcf5 Compare June 22, 2022 05:16

eickenberg commented Jun 23, 2022

View reviewed changes

MuawizChaudhary approved these changes Jun 23, 2022

View reviewed changes

janden approved these changes Jun 23, 2022

View reviewed changes

MAINT make sure that we check for numpy ndarrays in addition to numpy…

a7ec952

…-like ndarrays

lostanlen self-requested a review June 23, 2022 15:53

janden merged commit ee5fdcc into kymatio:dev Jun 23, 2022

lostanlen approved these changes Jun 23, 2022

View reviewed changes

		assert np.allclose(jax_result, numpy_result)


		def test_jnp_multiplication():

Uh oh!

Conversation

eickenberg commented Jun 30, 2021

Uh oh!

janden left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

janden Jul 2, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eickenberg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

eickenberg commented Jul 1, 2021

Uh oh!

eickenberg commented Jul 1, 2021

Uh oh!

janden commented Jul 2, 2021

Uh oh!

eickenberg commented Feb 14, 2022

Uh oh!

MuawizChaudhary commented May 1, 2022

Uh oh!

lostanlen commented May 2, 2022

Uh oh!

MuawizChaudhary left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

janden Jul 2, 2021 •

edited

Loading

MuawizChaudhary commented Jun 19, 2022 •

edited

Loading

lostanlen commented Jun 21, 2022 •

edited

Loading