gh-438: Add grad tests for all applicable functions by connoraird · Pull Request #1251 · glass-dev/glass

connoraird · 2026-10-08T16:36:36Z

Description

Adds gradients tests for all performance critical functions which take in and return an array.

Closes: #438

Changelog entry

Changed: Pinned JAX to <0.12.0 due to a change in multinomial behaviour.

Checks

Is your code passing linting?
Is your code passing tests?
Have you added additional tests (if required)?
Have you modified/extended the documentation (if required)?
Have you added a one-liner changelog entry above (if required)?

paddyroddy · 2026-10-19T08:51:31Z

I think the multiple dtype=jnp.float64 were needed otherwise the values default to integer which makes the functions not auto-differentiable.

Okay, I'll have a look. Wonder if we can force it with, e.g. 1.0.

paddyroddy

My comments are stale. I've done various changes to keep the formatting consistent with the rest of the codebase. This is a great piece of work!

I also fixed the readthedocs bug and a JAX bug coming from 0.12.0.

paddyroddy · 2026-10-16T18:00:46Z

    "ANN202", # missing-return-type-private-function
    "D100", # undocumented-public-module
    "D103", # TODO: undocumented-public-function
+    "FBT001", # boolean-type-hint-positional-argument


Probably makes sense

paddyroddy · 2026-10-16T18:02:54Z

+    jnp: ModuleType,
+) -> None:
+    """Tests glass.VolumeWeight is auto differentiable."""
+    z = jnp.linspace(0.1, 1, 6, dtype=jnp.float64)


Shouldn't need to specify the dtype here as we're using the jnp fixture.

paddyroddy · 2026-10-16T18:15:34Z

+
+pytest.importorskip("jax", reason="tests require jax")
+
+import jax.numpy as jnp


What about the fixture?

I had been expecting this as a series of PRs. Impressive work as one. But things like this would have been caught earlier.

Yes, in hind-sight multiple PRs would have been better.

connoraird

A few minor comments but I'm happy with it, if you are.

connoraird · 2026-10-19T15:43:12Z

    jnp: ModuleType,
-    size_input: int | tuple[int, ...] | None,
    shape_output: tuple[int, ...],
+    size_input: int | tuple[int, ...] | None,


There are a lot of changes unrelated to this PR in here. Perhaps that's okay since we're nearing the end but just wanted to mention it.

Totally fair to call out. I just spotted we had stopped following established patterns.

Got a little carried away.

paddyroddy · 2026-10-20T08:02:46Z

+  jobs:
+    post_checkout:
+      - git fetch --unshallow --tags origin || true


Fixes the readthedocs bug

paddyroddy · 2026-10-20T08:21:51Z

+def test_apply_visibility(jnp: ModuleType) -> None:
+    """Tests glass.points._apply_visibility is auto differentiable."""
+    k = (1, 1)
+    n = 24751.77674965 * jnp.ones(12)


This is very precise!

paddyroddy · 2026-10-20T08:29:40Z

This looks good now, going to merge!

connoraird self-assigned this Oct 8, 2026

connoraird added array-api Work is related to the Array API jax-porting Removing mutation to make the code JAX compatible labels Oct 8, 2026

connoraird force-pushed the connor/add-more-grad-tests branch from 2c91928 to deff382 Compare October 9, 2026 16:39

connoraird changed the title ~~gh-: Add grad tests for all applicable functions~~ gh-438: Add grad tests for all applicable functions Oct 9, 2026

connoraird mentioned this pull request Oct 15, 2026

gh-1152: Test any xp with cosmo fixture #1257

Merged

3 of 5 tasks

connoraird force-pushed the connor/add-more-grad-tests branch from d8891a8 to fb2f0aa Compare October 16, 2026 08:48

connoraird mentioned this pull request Oct 16, 2026

Fix type casting for auto-differentiability #1267

Open

connoraird force-pushed the connor/add-more-grad-tests branch from d497974 to 6b0c0f0 Compare October 16, 2026 13:24

connoraird changed the base branch from main to connor/1267-make-radwin-diffable October 16, 2026 13:24

connoraird changed the base branch from connor/1267-make-radwin-diffable to main October 16, 2026 14:48

connoraird force-pushed the connor/add-more-grad-tests branch from 6b0c0f0 to be78ea4 Compare October 16, 2026 14:49

connoraird added 18 commits October 16, 2026 18:02

Add grad tests for fields.py

e8e2977

Ignore FBT001 in tests

4a7833d

Add grad tests for lensing.py

971c311

Add grad tests for arraytools.py

ab0e72e

Add grad tests for galaxies.py

3055395

Add grad tests for harmonics.py

28d479d

Add grad tests for points.py

fd2bd5a

WIP: Add grad tests for shapes.py

63ba4cb

Add notes about issues with auto-differentiability

e008607

Add grad tests for shapes.py

70c31a3

Update test_galaxies to allow more auto differentiabilty

2722a70

Document issue with using 0.0 as an input

4c224cc

Fix test_shapes

3d356cb

Add grad tests for shells.py

3764f7b

Skip failing tests

79fe6e9

Add test with period and explain reasoning

7aaefc9

Fix some skipped tests in test_fields.py

d6f6261

Fix some skipped tests in test_arraytools.py

7df2676

paddyroddy added 18 commits October 19, 2026 11:40

Sort testing imports

422117d

Use ARRAY_BACKEND environment variable

9f97970

Add missing environ

25664bb

Remove env=os.environ

14d8f98

Minimise to only when test skip for JAX

3cb7975

Remove from imports

b707688

Fix linting

fe747a2

Refine array backend tests

12a5791

Remove 64bit from zeros

aece28b

Remove 64 bit from ones

0371004

Remove type from linspace

14d3dd8

Remove float64

d97ea3f

Sort all test inputs

c09ff7d

Order parameters

4d10542

Pin JAX

cd5f2a0

Messed up version

7d3c8cb

Pin JAX everywhere

9893756

Fix rtd

e70e22e

paddyroddy approved these changes Oct 19, 2026

View reviewed changes

paddyroddy mentioned this pull request Oct 19, 2026

gh-1269: Fix JAX multinomial for 0.12.0 #1270

Merged

3 of 5 tasks

connoraird commented Oct 19, 2026

View reviewed changes

paddyroddy added 2 commits October 20, 2026 08:50

Change backend test arrangement

e2b633b

Move importlib.util.find_spec into helper function

c8a4590

paddyroddy reviewed Oct 20, 2026

View reviewed changes

One line functions

2e42048

paddyroddy reviewed Oct 20, 2026

View reviewed changes

paddyroddy merged commit 8ffa463 into main Oct 20, 2026
15 of 16 checks passed

paddyroddy deleted the connor/add-more-grad-tests branch October 20, 2026 08:29


		pytest.importorskip("jax", reason="tests require jax")

		import jax.numpy as jnp

Conversation

connoraird commented Oct 8, 2026 • edited by paddyroddy Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changelog entry

Checks

Uh oh!

paddyroddy commented Oct 19, 2026

Uh oh!

paddyroddy left a comment

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!

connoraird 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

paddyroddy commented Oct 20, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

connoraird commented Oct 8, 2026 •

edited by paddyroddy

Loading