@@ -625,16 +625,16 @@ def _make_haloupdate(self, f, hse, key, sendrecv, **kwargs):
625625 lsizes , lofs = mapper [(d , LEFT , OWNED )]
626626 rsizes , rofs = mapper [(d , RIGHT , HALO )]
627627 args = [f , lsizes , lofs , rofs , rpeer , lpeer , comm ]
628- kwargs [ 'haloid' ] = len (body )
629- body . append ( self . _call_sendrecv ( sendrecv . name , * args , ** kwargs ))
628+ body . append ( self . _call_sendrecv ( sendrecv . name , * args , haloid = len (body ),
629+ ** kwargs ))
630630
631631 if (d , RIGHT ) in hse .halos :
632632 # Sending to right, receiving from left
633633 rsizes , rofs = mapper [(d , RIGHT , OWNED )]
634634 lsizes , lofs = mapper [(d , LEFT , HALO )]
635635 args = [f , rsizes , rofs , lofs , lpeer , rpeer , comm ]
636- kwargs [ 'haloid' ] = len (body )
637- body . append ( self . _call_sendrecv ( sendrecv . name , * args , ** kwargs ))
636+ body . append ( self . _call_sendrecv ( sendrecv . name , * args , haloid = len (body ),
637+ ** kwargs ))
638638
639639 iet = List (body = body )
640640
@@ -1305,7 +1305,6 @@ def _allocate_buffers(self, f, shape, entry):
13051305 # The `memfree_args` will be used to deallocate the buffer upon
13061306 # returning from C-land
13071307 self ._memfree_args .extend ([bufg_memfree_args , bufs_memfree_args ])
1308- return
13091308
13101309 def _arg_defaults (self , allocator , alias , args = None ):
13111310 # Lazy initialization if `allocator` is necessary as the `allocator`
@@ -1416,15 +1415,14 @@ def _arg_defaults(self, allocator, alias, args=None):
14161415 for d1 in f .dimensions :
14171416 if d1 in fixed :
14181417 continue
1418+ if d0 is d1 :
1419+ if region is OWNED :
1420+ sizes .append (getattr (f ._size_owned [d0 ], side .name ))
1421+ elif region is HALO :
1422+ sizes .append (getattr (f ._size_halo [d0 ], side .name ))
14191423 else :
1420- if d0 is d1 :
1421- if region is OWNED :
1422- sizes .append (getattr (f ._size_owned [d0 ], side .name ))
1423- elif region is HALO :
1424- sizes .append (getattr (f ._size_halo [d0 ], side .name ))
1425- else :
1426- sizes .append (self ._as_number (f ._size_nopad [d1 ], args ))
1427- mapper [(d0 , side , region )] = (sizes )
1424+ sizes .append (self ._as_number (f ._size_nopad [d1 ], args ))
1425+ mapper [(d0 , side , region )] = sizes
14281426
14291427 i = 0
14301428 for d in f .dimensions :
0 commit comments