WIP: ENH: Expand gufunc signature to allow flexible dimension specs by mattip · Pull Request #11132 · numpy/numpy

mattip · 2018-05-22T00:13:54Z

The discussion of issue #9028 yielded two different possible directions for allowing ndarray subclasses to override __matmul__ via __array_ufunc__. This is option 1 - extend the core signature of a generalized ufunc to allow a single ufunc matmul which will implement vector-vector, vector-matrix, matrix-vector, and matrix-matrix multiplication. In this PR, the syntax of a signature is expanded to allow (n?, k),(k, m?)->(n?, m?) syntax, where the ? means "this dimension is optional" [0],[1]. I also incorporated PR #5015 where signatures can have a fixed dimension (3),(3)->(3).

What changed?

There is a change in the PyUFuncObject structure. I added two fields, and repurposed the unused field reserved1 to become a version field for downstream consumers of the C-API.
The parsing code during ufunc instantiation now parses the expanded signature
When calling a ufunc with flexible or fixed dimensions, extra checks are done to ensure consistency

What is yet to be done?

The code needs cleanup and refactoring
Tests were added but there is a good chance I missed corner cases
I wrote the code together with the PR making matmul a true ufunc, I probably have chages there that belong here and visa-versa

What is next?

Use this mechanism to make matmul a true ufunc
Get feedback on the relative merits of this approach vs. the one taken in PR WIP: ENH: allow __array_ufunc__ to override __matmul__ #11061, where I produced a demi-ufunc with a wrapper function that can be used to convert any function into a ufunc.

[0]: Note that, following the PEP for matrix multiplication, there is no broadcasting over the optional dimensions, in the example given if the first operand has 2 or more dimensions, the second-to-the-last will be 'n', no newaxis will be inserted instead.

[1] The syntax '(n,k),(k,m)->(n,m);(n,k),(k)->(n);(k),(k,m)->(m);(k),(k)->()` was also considered and rejected as allowing too much flexibility in stacking alternate core signatures.

mhvk

@mattip - this is quite wonderful! I did find a few possible problems. It would also be good to add tests actually using the cross product (I guess matmul will do for the flexible ones!)

mhvk · 2018-05-23T17:31:04Z

    ufunc->core_offsets = PyArray_malloc(sizeof(int) * ufunc->nargs);
-    if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL
-        || ufunc->core_offsets == NULL) {
+    /* The next two items will be shrunk later */


removing szs makes it two again, so "fixed"

mhvk · 2018-05-23T17:45:37Z

+                parse_error = "expect dimension name or non-zero frozen size";
                goto fail;
            }
            while (j < ufunc->core_num_dim_ix) {


This while loop is not necessarily correct for frozen dimensions, since a number does not have to be given exactly the same way (say, a signature of (03)->(3). How about the following:

for (j = 0; j < ufunc->core_num_dim_ix; j++) { if (frozen_size >= 0 ? frozen_size == ufunc->core_dim_szs[j]) : _is_same_name(signature+i, var_names[j])) { break; } }

(Could also be a while, but then at least include the test for (in)equality in the condition!)

Though perhaps it is better to just split the frozen and normal paths altogether (see above)

fixed by removing code

mhvk · 2018-05-23T18:04:03Z

+                goto fail;
+            }
+
+            if (!_is_alpha_underscore(signature[i]) &&


I think here and below it might make more sense to separate out more:

if (_is_alpha...) { check on var_names } else { frozen_size = _get_... if (frozen_size < 0) { parse error } check on same frozen size }

fixed by removing frozen size code

mhvk · 2018-05-23T18:09:21Z

 */
 static int
 _get_coredim_sizes(PyUFuncObject *ufunc, PyArrayObject **op,
+                   int *core_num_dims, int * flexible_activated,


extraneous space after *

fixed by removing function

mhvk · 2018-05-23T18:10:23Z

    for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
-        core_dim_sizes[i] = -1;
+        /* support fixed-size dim names */
+        core_dim_sizes[i] = ufunc->core_dim_szs[i];


Maybe call the ufunc part core_dim_sizes as well?

fixed by removing ufunc->core_dim_szs

mhvk · 2018-05-23T18:37:27Z

        if (i >= nin) {
            int dim_offset = ufunc->core_offsets[i];
            int num_dims = core_num_dims[i];
+            int has_flexible=0;


spaces around =

Also, can we call this num_flexible, since that is what it ends up being?

mhvk · 2018-05-23T18:38:15Z

+                    int core_index = ufunc->core_dim_ixs[dim_offset + idim];
+                    if (flexible_activated[core_index]) {
+                        /* skip it */
+                        has_flexible ++;


no space before ++

mhvk · 2018-05-23T18:40:56Z

    ufunc->core_dim_ixs = NULL;
    ufunc->core_offsets = NULL;
    ufunc->core_signature = NULL;
+    ufunc->core_dim_szs = NULL;


Should also initialize core_dim_flexible

mhvk · 2018-05-23T18:41:16Z

 {
    PyArray_free(ufunc->core_num_dims);
    PyArray_free(ufunc->core_dim_ixs);
+    PyArray_free(ufunc->core_dim_szs);


Should also free ->core_dim_flexible or we'll have a memory leak...

mhvk · 2018-05-23T18:49:45Z

+        assert_equal(ixs, (0, 1, 1, 2, 0, 2))
+        assert_equal(szs, (-1, -1, -1))
+        assert_equal(flex, (1, 0, 1))



Can you add tests for the fixed sizes too? Including travesties like (+0003)->(3)...

removing fixed sizes to a different pull request

"fixed" by removing fixed sizes

mattip · 2018-05-23T19:47:32Z

I am waiting to proceed until there is consensus around the general validity of the approach. How do you feel about changing the signature and PyUFuncObject c-compatibility vs. PR #11061 with the wrapper function?

mhvk · 2018-05-24T01:28:52Z

I'd really like the fixed dimensions independently of matmul.

I'm slightly less sure about the ? -- it seems matmul specific. But a possible follow-up which maybe allows combining p? with p, which would indicate allowing the marked one to be broadcast (ie., can be 1) would I think be quite useful.

In any case, I do think I like this better than the multiple signatures idea.

mhvk · 2018-05-24T01:32:46Z

More to the point: for matmul at least, I like this better as it is easier for me to deal with: in __array_ufunc__, I can just check that ufunc is np.matmul and act accordingly. It is still much less clear how to do that for the wrapper.

eric-wieser · 2018-05-24T01:58:45Z

I haven't looked over this in detail, but...

If we're trying to avoid a scenario where matmul is a wrapper that dispatches to 1 of a set of 3 ufuncs based on .shape, then we also should solve that problem for many of the functions in linalg which currently do the same.:

Let's not add special cases unless they also work for these cases too.

mattip · 2018-05-24T02:14:09Z

It is still much less clear how to do that for the wrapper.

ufunc.__name__. I would think that if you are comparing to np.matmul or np.add, you are setting yourself up for trouble when you are given a different ndarray subclass that has overridden the np.ufunc

mhvk · 2018-05-24T13:58:14Z

ufunc.__name__ may well have problems with namespaces (which are a honking great idea!). I don't understand how another ndarray subclass could override the np.ufunc; you mean that their __array_ufunc__ decides to pass on something else than the ufunc they get passed in? If so, it is fine if Quantity fails - it doesn't know what to do with that ufunc.

mhvk · 2018-05-24T13:59:08Z

p.s. If the p? description is controversial, I'd be in favour of still doing the fixed dimensions.

mhvk · 2018-05-25T14:34:22Z

+        /*
+         * for each core_num_dim_ix, 1 for flexible (signature has ?) 0 otherwise
+         */
+        int *core_dim_flexible;


Implementation suggestion: I think we might as well make this a set of flags, so that we can re-use them for future expansions.

In fact, one could then add not just a flag UFUNC_COREDIM_FLEXIBLE but also a UFUNC_COREDIM_FIXED - for that case, ~~core_num_dim_idx~~ core_dim_ixs could just be set to the size, since there is no need to track the index.

I'm going ahead and try this - I guess one can see it as my contribution to the BIDS sprint!

mattip · 2018-05-26T01:38:58Z

we also should solve that problem for many of the functions in linalg which currently do the same.:

solve

svd

lstsq

solve: special handling of 1-less-ndim on inputs into a ufunc
svd_m and svd_n: choice is made dependent on m<n for a.shape[-2] == (m,n)
lstsq_m and lstsq_n choice is made as in svd

It is not clear to me how to refactor these to use the ? notation, but I am open to suggestions.

mattip · 2018-05-26T01:52:03Z

removed fix-dimension gufunc signatures and fixed review comments

mhvk · 2018-05-26T05:39:07Z

@mattip, @eric-wieser - I should really have read all the comments up here, but didn't have time tonight, yet still was thinking about all this. Some specific thought:

For ?, I think there are more things we can solve with this. Eg., there is the long-standing all_equal PR, which currently has signature (n),(n)->(). But really one would like to do all_equal(a, const, axis=...), which is not currently possible. Now suppose we used (n?),(n?)->() and interpreted the ? as follows: a dimension can be broadcast against others with the same character (with usual rule that if absent, it is broadcast automatically). That would solve all_equal, but it means the information should be kept per core dim, not per "distinct core dim". If we go that route, then (n?),(n)->(n) would mean that the first argument can be broadcast against the second, but the second just decides n. For the output, I think the question mark is similar if output is given, but if it is not then by default n? should remove the axis if n=1 while it should keep it regardless if one uses n. So, for the matmul ufunc, with signature (p?,n),(n, q?)->(p?, q?), it means that if one passes in a "vector" with shape (...,1,n) to multiply with a matrix (..., n, q), then the output would be (..., q) (this would mean matmul could do more than @, which is fine).
For fixed dimensions, I'm similarly wondering if one shouldn't just have an array with actual core sizes, not the distinct ones, i.e., when I pass in a number, that component dimension just gets set to that number, and the index to the common distinct core dimension ignored. Though perhaps this makes it too complicated with what the inner loop gets passed.
For the lsq routines, I think one effectively needs a calculation to decide on the output shape; one solution might be to provide a default output shape calculator, which can be overridden (just like we have a ufunc->type_resolver). Alternatively, in the context of ?, what about allowing, for outputs, a signature consisting only of ? (say (n),(m,n)->(?))? We could then not raise any errors until the ufunc->type_resolver is called, and then the user can override that one (I use that override currently to implement fixed-sized inputs...).

mhvk · 2018-05-26T15:03:54Z

@mattip - I am some way to the per-core (for each operand) flags (well, tests pass again).

mhvk · 2018-05-26T17:37:42Z

And am realizing that my description above is pretty poor - we're not really about broadcasting here (size=1), but rather about ignoring (size=0).

mhvk · 2018-05-28T14:25:08Z

@mattip - as you can see in #11175, I moved on from your earlier version which included the frozen dimensions. I think all of this is a great idea, and should work. However, the PRs are becoming quite big, and obviously with your reverting of the frozen dimensions, we've also diverged. I think it may be best to proceed in smaller steps, and took the liberty of going ahead with the first two, which seemed uncontroversial (I tried to cherry-pick, etc., so you get appropriate credit):

Your very nice clarification of the docs could be a quick separate PR - I went ahead and cherry-picked the relevant parts from this PR; see DOC: improvement of the documentation for gufunc. #11177.
Same for your expansion of the signature testing, first just with current master; see TST: Test dimensions/indices found from parsed gufunc signatures. #11178.
Then maybe frozen sizes as it is less intrusive.
Then add flags and have the flexible dimension.
Add broadcasting support.

For flexible in particular, my main question is whether you think going the route of flags is a good idea. Implementation-wise, it doesn't matter so much, but it gives more freedom.

On the implementation: #11175 ends up being quite a bit simpler (I think) than the proof of concept here. But again it may be useful to proceed in smaller steps, e.g., first doing the parser and then the consequence.

mattip · 2018-05-29T18:11:42Z

The flag proposal looks better. Closing this in favor of PR #11175

mhvk · 2018-05-29T19:04:23Z

OK, happy you like it. What I think I'll do is to squash our respective commits in #11175, and then rebase to get rid of the merge conflicts.

mattip added 3 commits May 18, 2018 13:58

ENH: Add frozen dimensions to gufunc signatures

e9a5ab5

convert reserved field to version, set new version to 1

cbbe0f6

add tests for new signatures

dbbeb58

mattip changed the title ~~ENH: Expand gufunc signature to allow flexible dimension specs~~ WIP: ENH: Expand gufunc signature to allow flexible dimension specs May 22, 2018

mattip mentioned this pull request May 22, 2018

WIP: ENH: Matmul as ufunc #11133

Closed

mattip added 2 commits May 22, 2018 14:23

test, document, implement flexible signature

18505d7

improve docs for PyUFunc_FromFuncAndData

67e752e

mattip force-pushed the expand-gufunc-signature branch 2 times, most recently from c2afb3c to b785515 Compare May 23, 2018 00:35

charris added 01 - Enhancement 15 - Discussion labels May 23, 2018

mhvk mentioned this pull request May 23, 2018

Wrap erfa functions with ufuncs astropy/astropy#7484

Closed

mhvk reviewed May 23, 2018

View reviewed changes

mhvk reviewed May 25, 2018

View reviewed changes

mattip added 3 commits May 25, 2018 18:49

remove fixed dimension ufuncsignatures

50d59fd

belongs on matmul-as-ufunc branch

d479a5b

fixes from review

f3d1622

mattip force-pushed the expand-gufunc-signature branch from b785515 to f3d1622 Compare May 26, 2018 01:50

mattip closed this May 29, 2018

mattip deleted the expand-gufunc-signature branch October 9, 2018 17:20

Conversation

mattip commented May 22, 2018

Uh oh!

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

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattip commented May 23, 2018

Uh oh!

mhvk commented May 24, 2018

Uh oh!

mhvk commented May 24, 2018

Uh oh!

eric-wieser commented May 24, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mattip commented May 24, 2018

Uh oh!

mhvk commented May 24, 2018

Uh oh!

mhvk commented May 24, 2018

Uh oh!

mhvk May 25, 2018 • 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!

mattip commented May 26, 2018

Uh oh!

mattip commented May 26, 2018

Uh oh!

mhvk commented May 26, 2018

Uh oh!

mhvk commented May 26, 2018

Uh oh!

mhvk commented May 26, 2018

Uh oh!

mhvk commented May 28, 2018

eric-wieser commented May 24, 2018 •

edited

Loading

mhvk May 25, 2018 •

edited

Loading