So the implementation from intel would replace the axis_mul_or_truediv function we use as last step here:

scanpy/src/scanpy/preprocessing/_normalization.py

Lines 28 to 52 in 834159a

    
           def _normalize_data(X, counts, after=None, *, copy: bool = False): 
        
               X = X.copy() if copy else X 
        
               if issubclass(X.dtype.type, int | np.integer): 
        
                   X = X.astype(np.float32)  # TODO: Check if float64 should be used 
        
               if after is None: 
        
                   if isinstance(counts, DaskArray): 
        
                       def nonzero_median(x): 
        
                           return np.ma.median(np.ma.masked_array(x, x == 0)).item() 
        
                       after = da.from_delayed( 
        
                           dask.delayed(nonzero_median)(counts), 
        
                           shape=(), 
        
                           meta=counts._meta, 
        
                           dtype=counts.dtype, 
        
                       ) 
        
                   else: 
        
                       counts_greater_than_zero = counts[counts > 0] 
        
                       after = np.median(counts_greater_than_zero, axis=0) 
        
               counts = counts / after 
        
               out = X if isinstance(X, np.ndarray | CSBase) else None 
        
               return axis_mul_or_truediv( 
        
                   X, counts, op=truediv, out=out, allow_divide_by_zero=False, axis=0 
        
               )

Because everything done until that point is to find median in case there is no fixed target sum given. So I would essentially have to utilize numba in axis_mul_or_truediv, does that sound right @ilan-gold @flying-sheep ? If so I would proceed on that

	def _normalize_data(X, counts, after=None, *, copy: bool = False):
	X = X.copy() if copy else X
	if issubclass(X.dtype.type, int \| np.integer):
	X = X.astype(np.float32) # TODO: Check if float64 should be used
	if after is None:
	if isinstance(counts, DaskArray):

	def nonzero_median(x):
	return np.ma.median(np.ma.masked_array(x, x == 0)).item()

	after = da.from_delayed(
	dask.delayed(nonzero_median)(counts),
	shape=(),
	meta=counts._meta,
	dtype=counts.dtype,
	)
	else:
	counts_greater_than_zero = counts[counts > 0]
	after = np.median(counts_greater_than_zero, axis=0)
	counts = counts / after
	out = X if isinstance(X, np.ndarray \| CSBase) else None
	return axis_mul_or_truediv(
	X, counts, op=truediv, out=out, allow_divide_by_zero=False, axis=0
	)

normalize_total with numba #3135

Description

What kind of feature would you like to request?

Please describe your wishes

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`normalize_total` with numba #3135