IPython profile to automatically initialize PyCUDA on a specified device

ipy_profile_pycuda.py

#!/usr/bin/env python

"""
IPython 0.10 profile to automatically initialize PyCUDA on a specified
device.
"""

import sys
import atexit
import IPython.ipapi
ip = IPython.ipapi.get()

ip.ex("import numpy as np")
ip.ex("import atexit")
ip.ex("import pycuda.driver as drv")
ip.ex("drv.init()")

dev_num_str = sys.argv[-1] if sys.argv[-1].isdigit() else '0'
ip.ex("dev = drv.Device(%s)" % dev_num_str)
ip.ex("atexit.register(dev.make_context().pop)")
ip.ex("import pycuda.gpuarray as gpuarray")

ip.ex("print ")
ip.ex("print '*** pycuda initialized on ' + dev.name() + ' [%s] ***'" % dev_num_str)
ip.ex("print ")