def forward(ctx, formula, aliases, backend, dtype, device_id, ranges, accuracy_flags, *args): optional_flags = ['-DPYTORCH_INCLUDE_DIR=' + ';'.join(include_dirs)] + accuracy_flags myconv = LoadKeOps(formula, aliases, dtype, 'torch', optional_flags).import_module() # Context variables: save everything to compute the gradient: ctx.formula = formula ctx.aliases = aliases ctx.backend = backend ctx.dtype = dtype ctx.device_id = device_id ctx.ranges = ranges ctx.accuracy_flags = accuracy_flags ctx.myconv = myconv tagCPUGPU, tag1D2D, tagHostDevice = get_tag_backend(backend, args) if tagCPUGPU==1 & tagHostDevice==1: device_id = args[0].device.index for i in range(1,len(args)): if args[i].device.index != device_id: raise ValueError("[KeOps] Input arrays must be all located on the same device.") if ranges is None : ranges = () # To keep the same type result = myconv.genred_pytorch(tagCPUGPU, tag1D2D, tagHostDevice, device_id, ranges, *args) # relying on the 'ctx.saved_variables' attribute is necessary if you want to be able to differentiate the output # of the backward once again. It helps pytorch to keep track of 'who is who'. ctx.save_for_backward(*args, result) return result
def forward(ctx, formula, aliases, backend, dtype, device_id, ranges, optional_flags, rec_multVar_highdim, nx, ny, *args): # N.B. when rec_multVar_highdim option is set, it means that formula is of the form "sum(F*b)", where b is a variable # with large dimension. In this case we set compiler option MULT_VAR_HIGHDIM to allow for the use of the special "final chunk" computation # mode. However, this may not be also true for the gradients of the same formula. In fact only the gradient # with respect to variable b will have the same form. Hence, we save optional_flags current status into ctx, # before adding the MULT_VAR_HIGHDIM compiler option. ctx.optional_flags = optional_flags.copy() if rec_multVar_highdim is not None: optional_flags += ["-DMULT_VAR_HIGHDIM=1"] myconv = LoadKeOps(formula, aliases, dtype, 'torch', optional_flags, include_dirs).import_module() # Context variables: save everything to compute the gradient: ctx.formula = formula ctx.aliases = aliases ctx.backend = backend ctx.dtype = dtype ctx.device_id = device_id ctx.ranges = ranges ctx.rec_multVar_highdim = rec_multVar_highdim ctx.myconv = myconv ctx.nx = nx ctx.ny = ny tagCPUGPU, tag1D2D, tagHostDevice = get_tag_backend(backend, args) if tagCPUGPU == 1 & tagHostDevice == 1: device_id = args[0].device.index for i in range(1, len(args)): if args[i].device.index != device_id: raise ValueError( "[KeOps] Input arrays must be all located on the same device." ) if ranges is None: ranges = () # To keep the same type # N.B.: KeOps C++ expects contiguous integer arrays as ranges ranges = tuple(r.contiguous() for r in ranges) result = myconv.genred_pytorch(tagCPUGPU, tag1D2D, tagHostDevice, device_id, ranges, nx, ny, *args) # relying on the 'ctx.saved_variables' attribute is necessary if you want to be able to differentiate the output # of the backward once again. It helps pytorch to keep track of 'who is who'. ctx.save_for_backward(*args, result) return result