gh-1021: Add decorator for np conversions by prady0t · Pull Request #1137 · glass-dev/glass

prady0t · 2026-06-16T12:06:54Z

Description

Addresses #1021

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)?

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

prady0t · 2026-06-23T12:40:41Z

        dxp = default_xp(xp.__name__)
        return tuple(xp.asarray(arr) for arr in dxp.tril_indices(n, k=k, m=m))
+
+def numpy_fallback(func):


@lucascolley This could be more generalised for the standard. Currently, I feel the limitation is the inability to specify parameters such as dtype for xp.asarray.

@ntessore @paddyroddy Would love some feedback.

cc @betatim, I wonder whether there is any overlap between this and scikit-learn/scikit-learn#34324 which would be nice to upstream into xpx.numpy.

This is pretty much out of my purview, and @paddyroddy is on leave currently, so it might be a bit before he can answer.

In my opinion, which as I said doesn't carry much weight here, I don't think this is something where GLASS should have to innovate. We're in the cosmological simulations business, not the array business 🙂

There isn't a lot of context in the PR description, what I figure from a quick browse: this PR adds a decorator that converts all "arrays" that are passed to a function to numpy arrays, then calls the wrapped function and converts any "arrays" in the output back to the original namespace.

I'm not sure how much the work in scikit-learn/scikit-learn#34324 is useful for others. It feels somewhat scikit-learn specific. The tricky (imho) part is deciding when to convert to numpy, which I guess is a library specific thing? At least in scikit-learn the decision to convert to numpy is not a simple "this function is hard to support". It depends on the global config, the specific constructor arguments of the class and what the type of the input array is. Once you know you want to convert to numpy it feels like you are "home free".

For this library it looks like most of the code of the decorator deals with finding arrays in the input arguments and such infrastructure. One thing I'm wondering is what to do if the inputs use different array namespaces or if one of the inputs is a torch array and another a list. But maybe these are things aren't an issue here? In scikit-learn we allow a mix of things :-/

Conclusion, I wonder if we can make a generic fallback decorator for xpx that doesn't get mega complex and at the same time supports most of the different behaviours users would want. If the answer is "yes we can" then I think it would be a useful tool (even if I don't think we'd use it in scikit-learn).

paddyroddy

This seems like a really nice approach and is the kind of thing I was thinking of. Long-term I definitely wouldn't want the code in this repo as it is complex, and our current approach would be easier for contributors to understand. Do we think we've captured all edge cases here? I assume the decorator is smart enough to detect which variables need to be converted and which don't?

paddyroddy · 2026-07-02T15:17:38Z

    # This function is difficult to port to the Array API so for now we work
    # in NumPy and ultimately convert back at the end of it.


We would need to remove this comment too

paddyroddy · 2026-07-02T15:19:30Z

        dxp = default_xp(xp.__name__)
        return tuple(xp.asarray(arr) for arr in dxp.tril_indices(n, k=k, m=m))
+
+def numpy_fallback(func):


Some typing and a docstring would be good here

Thanks for the review. I've added some in the latest commit.

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

…/prady0t/glass into add-decorator-for-np-conversions

prady0t

Currently, the state of the decorator is:

Goes through all the function parameters (args and kwargs) and collects those which have __array_namespace__ attribute.
If xp is in kwargs, it uses its value. If value is none, then use default_xp(). xp value overwrites parameter's __array_namespace__ attribute value.
If no xp parameter and no parameter has __array_namespace__ attribute, returns the function as is.
Parameters are converted to numpy
Depending on what the original function returns, the wrapper recursively converts (so that it can support tuple, list and dict) back to the xp.

There are two places where we cannot use this decorator: pixwin and the query_strip function, as we are specifying dtype during conversion. Something that this decorator cannot do yet.

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

prady0t · 2026-07-08T21:03:19Z

+        def collect_arrays(obj):
+            if hasattr(obj, "__array_namespace__"):
+                list_of_xp.append(obj)
+            elif isinstance(obj, dict):
+                for v in obj.values():
+                    collect_arrays(v)
+            elif isinstance(obj, Sequence) and not isinstance(obj, (str, bytes)):
+                for item in obj:
+                    collect_arrays(item)


I've made the decorator even smarter! Now, it's able to extract xp if the given arguments are a collection (or Sequence) of xp. Similarly, NumPy conversions are also being done recursively. This is helpful in the places commented below.

@paddyroddy gentle ping for a review.

Sorry I've been on leave

prady0t · 2026-07-08T21:04:13Z

+    return healpy.alm2map(
+        alms,
+        nside,
+        inplace=inplace,
+        lmax=lmax,
+        pixwin=pixwin,
+        pol=pol,


prady0t · 2026-07-08T21:10:09Z

+    return healpy.map2alm(
+        maps,
+        datapath=_get_healpy_datapath(),
+        lmax=lmax,
+        pol=pol,
+        use_pixel_weights=use_pixel_weights,


prady0t · 2026-07-08T21:10:59Z

    theta, phi = healpix.randang(
        nside,
-        np.asarray(ipix),
+        ipix,


prady0t · 2026-07-08T21:11:12Z

-        return xp.asarray(
-            healpy.Rotator(coord=self.coord).rotate_map_pixel(np.asarray(m)),
-        )
+        return healpy.Rotator(coord=self.coord).rotate_map_pixel(m)


Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

prady0t and others added 3 commits June 16, 2026 17:27

Adding decorator for numpy fallback

3f7f192

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

lint

ec0aa32

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

Merge branch 'main' into add-decorator-for-np-conversions

6e9fabf

prady0t changed the title ~~Add decorator for np conversions~~ ENH: Add decorator for np conversions Jun 16, 2026

prady0t changed the title ~~ENH: Add decorator for np conversions~~ gh-1021: Add decorator for np conversions Jun 16, 2026

prady0t commented Jun 23, 2026

View reviewed changes

Merge branch 'main' into add-decorator-for-np-conversions

805df5e

paddyroddy reviewed Jul 2, 2026

View reviewed changes

prady0t added 2 commits July 6, 2026 15:31

Adding typing + docstring + suggestion

7eb5d89

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

Merge branch 'add-decorator-for-np-conversions' of https://github.com…

89c1f54

…/prady0t/glass into add-decorator-for-np-conversions

prady0t commented Jul 6, 2026

View reviewed changes

paddyroddy reviewed Jul 7, 2026

View reviewed changes

Comment thread glass/healpix.py

paddyroddy linked an issue Jul 8, 2026 that may be closed by this pull request

Decorator for input output conversion array-api porting #1021

Open

paddyroddy assigned prady0t Jul 8, 2026

paddyroddy added enhancement New feature or request maintenance Maintenance: refactoring, typos, etc. array-api Work is related to the Array API labels Jul 8, 2026

paddyroddy and others added 2 commits July 8, 2026 17:40

Merge branch 'main' into add-decorator-for-np-conversions

a717135

making the decorrator smarter

33b8a29

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

prady0t commented Jul 8, 2026

View reviewed changes

prady0t marked this pull request as ready for review July 8, 2026 21:12

paddyroddy reviewed Jul 14, 2026

View reviewed changes

Comment thread array-api-tests Outdated

remove accidently added array_api_tests folder

cf6c9e9

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

		# This function is difficult to port to the Array API so for now we work
		# in NumPy and ultimately convert back at the end of it.

Uh oh!

Conversation

prady0t commented Jun 16, 2026

Description

Checks

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lucascolley Jun 23, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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!

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

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

lucascolley Jun 23, 2026 •

edited

Loading