Run black

Author: aseyboldt

pytensor/link/numba/dispatch/elemwise.py

@@ -447,7 +447,7 @@ def _vectorize_bc(
     noalias_outputs=False,
 ):
 
-    flags = True 
+    flags = True
     {
         "arcp",  # Allow Reciprocal
         "contract",  # Allow floating-point contraction
@@ -491,7 +491,10 @@ def codegen(context, builder, signature, args):
             # Lower the code of the scalar function so that we can use it in the inner loop
             # Caching is set to false to avoid a numba bug TODO ref?
             inner_func = context.compile_subroutine(
-                builder, scalar_func, scalar_signature, caching=False,
+                builder,
+                scalar_func,
+                scalar_signature,
+                caching=False,
             )
             inner = inner_func.fndesc
 
@@ -508,12 +511,12 @@ def extract_array(aryty, ary):
             # TODO I think this is better than the noalias attribute
             # for the input, but self_ref isn't supported in a released
             # llvmlite version yet
-            #mod = builder.module
-            #domain = mod.add_metadata([], self_ref=True)
-            #input_scope = mod.add_metadata([domain], self_ref=True)
-            #output_scope = mod.add_metadata([domain], self_ref=True)
-            #input_scope_set = mod.add_metadata([input_scope, output_scope])
-            #output_scope_set = mod.add_metadata([input_scope, output_scope])
+            # mod = builder.module
+            # domain = mod.add_metadata([], self_ref=True)
+            # input_scope = mod.add_metadata([domain], self_ref=True)
+            # output_scope = mod.add_metadata([domain], self_ref=True)
+            # input_scope_set = mod.add_metadata([input_scope, output_scope])
+            # output_scope_set = mod.add_metadata([input_scope, output_scope])
 
             inputs = [
                 extract_array(aryty, ary)
@@ -559,8 +562,8 @@ def extract_array(aryty, ary):
                     context, builder, *array_info, idxs_bc, *safe
                 )
                 val = builder.load(ptr)
-                #val.set_metadata("alias.scope", input_scope_set)
-                #val.set_metadata("noalias", output_scope_set)
+                # val.set_metadata("alias.scope", input_scope_set)
+                # val.set_metadata("noalias", output_scope_set)
                 input_vals.append(val)
 
             # Call scalar function
@@ -581,13 +584,13 @@ def extract_array(aryty, ary):
             ):
                 if accu is not None:
                     load = builder.load(accu)
-                    #load.set_metadata("alias.scope", output_scope_set)
-                    #load.set_metadata("noalias", input_scope_set)
+                    # load.set_metadata("alias.scope", output_scope_set)
+                    # load.set_metadata("noalias", input_scope_set)
                     new_value = builder.fadd(load, value)
                     store = builder.store(new_value, accu)
                     # TODO ?
-                    #store.set_metadata("alias.scope", output_scope_set)
-                    #store.set_metadata("noalias", input_scope_set)
+                    # store.set_metadata("alias.scope", output_scope_set)
+                    # store.set_metadata("noalias", input_scope_set)
                 else:
                     idxs_bc = [
                         zero if bc else idx
@@ -596,10 +599,12 @@ def extract_array(aryty, ary):
                     ptr = cgutils.get_item_pointer2(
                         context, builder, *outputs[i], idxs_bc
                     )
-                    #store = builder.store(value, ptr)
-                    store = arrayobj.store_item(context, builder, out_types[i], value, ptr)
-                    #store.set_metadata("alias.scope", output_scope_set)
-                    #store.set_metadata("noalias", input_scope_set)
+                    # store = builder.store(value, ptr)
+                    store = arrayobj.store_item(
+                        context, builder, out_types[i], value, ptr
+                    )
+                    # store.set_metadata("alias.scope", output_scope_set)
+                    # store.set_metadata("noalias", input_scope_set)
 
             # Close the loops and write accumulator values to the output arrays
             for depth, loop in enumerate(loop_stack[::-1]):
@@ -615,12 +620,14 @@ def extract_array(aryty, ary):
                             context, builder, *outputs[output], idxs_bc
                         )
                         load = builder.load(accu)
-                        #load.set_metadata("alias.scope", output_scope_set)
-                        #load.set_metadata("noalias", input_scope_set)
-                        #store = builder.store(load, ptr)
-                        store = arrayobj.store_item(context, builder, out_types[output], load, ptr)
-                        #store.set_metadata("alias.scope", output_scope_set)
-                        #store.set_metadata("noalias", input_scope_set)
+                        # load.set_metadata("alias.scope", output_scope_set)
+                        # load.set_metadata("noalias", input_scope_set)
+                        # store = builder.store(load, ptr)
+                        store = arrayobj.store_item(
+                            context, builder, out_types[output], load, ptr
+                        )
+                        # store.set_metadata("alias.scope", output_scope_set)
+                        # store.set_metadata("noalias", input_scope_set)
                 loop.__exit__(None, None, None)
             return
 
@@ -672,6 +679,7 @@ def make_output(iter_shape, bc, dtype):
 
             check_arrays = check_broadcasting
         else:
+
             @numba.extending.register_jitable
             def make_output(iter_shape, bc, dtype):
                 return np.empty((), dtype)
@@ -680,7 +688,6 @@ def make_output(iter_shape, bc, dtype):
             def check_arrays(a, b, c):
                 pass
 
-
         make_outputs = tuple_mapper(make_output)
 
         def impl(*inputs):
@@ -756,10 +763,7 @@ def numba_funcify_Elemwise(op, node, **kwargs):
         scalar_op_fn,
         input_bc_patterns,
         output_bc_patterns,
-        output_dtypes=tuple([
-            variable.dtype
-            for variable in node.outputs
-        ]),
+        output_dtypes=tuple([variable.dtype for variable in node.outputs]),
     )
 
     # TODO We should do this in vectorize instead
@@ -771,13 +775,16 @@ def elemwise_inplace(*inputs):
             outputs = vectorized(*inputs)
             for out_idx, in_idx in literal_unroll(pattern):
                 inputs[in_idx][...] = outputs[out_idx]
+
     else:
         elemwise_inplace = vectorized
 
     if len(node.outputs) == 1:
+
         @numba_njit
         def elemwise_wrapper(*inputs):
             return elemwise_inplace(*inputs)[0]
+
     else:
         elemwise_wrapper = vectorized