clemisch

import jax
import jax.numpy as np
import numpy as onp


def slice_in_dim(operand, start_index, limit_index, stride=1, axis=0):
  """Convenience wrapper around slice applying to only one dimension."""
  start_indices = [0] * operand.ndim
  limit_indices = list(operand.shape)
  strides = [1] * operand.ndim

clemisch

import numpy as np
import itk
from itk import RTK as rtk

GPU_IMG = rtk.CudaImage[itk.F, 3]
CPU_IMG = rtk.Image[itk.F, 3]

def cpu_to_gpu_image(cpu_img, gpu_img=None):
    if gpu_img is None:
        gpu_img = GPU_IMG.New()