agd.AutomaticDifferentiation.cupy_support

 1import numpy as np
 2from .Base import implements_cupy_alt,cp,is_ad
 3
 4"""
 5This file implements a few numpy functions that not well supported by the
 6cupy version (6.0, thus outdated) that is available on windows by conda at the 
 7time of writing.
 8"""
 9
10# -- NOT -- fixed in cupy 8.2
11@implements_cupy_alt(np.max,TypeError)
12def max(a,*args,**kwargs):
13	initial=kwargs.pop('initial') # cupy (old version ?) does not accept initial argument
14	return np.maximum(initial,np.max(a,*args,**kwargs))
15
16def _along_axis(arr,indices,axis):
17	axis%=arr.ndim
18	def indices_(ax):
19		if ax==axis: return indices
20		sax = arr.shape[ax]
21		ind = np.arange(sax).reshape((1,)*ax + (sax,)+(1,)*(arr.ndim-ax-1))
22		return np.broadcast_to(ind,indices.shape)
23	return tuple(indices_(ax) for ax in range(arr.ndim))
24
25# --- NOT --- fixed in cupy 8.2
26@implements_cupy_alt(np.put_along_axis,TypeError)
27def put_along_axis(arr,indices,values,axis):
28	arr[_along_axis(arr,indices,axis)]=values
29
30# -- NOT -- fixed in cupy 8.2
31@implements_cupy_alt(np.packbits,TypeError)
32def packbits(arr,bitorder='big'):
33	"""Implements bitorder option in cupy""" 
34	if bitorder=='little':
35		shape = arr.shape
36		arr = arr.reshape(-1,8)
37		arr = arr[:,::-1]
38		arr = arr.reshape(shape)
39	return cp.packbits(arr)