Tensor plotting functions

plot.py

from matplotlib import pyplot as plt
from mpl_toolkits import mplot3d

def scatter(X, Y, c=None, ax=None):
  ax = plt.axes() if ax is None else ax
  ax.scatter(X.numpy(), Y.numpy(), c=None if c is None else c.numpy())
  return ax

def contour(X, Y, Z, levels=None, ax=None):
  ax = plt.axes() if ax is None else ax
  ax.contour(X.numpy(), Y.numpy(), Z.numpy(), levels=levels)
  return ax

def scatter3D(X, Y, Z, ax=None):
  ax = plt.axes(projection='3d') if ax is None else ax
  ax.scatter3D(X.numpy(), Y.numpy(), Z.numpy(), c=Z.numpy())
  return ax

def plot3D(X, Y, Z, ax=None):
  ax = plt.axes(projection='3d') if ax is None else ax
  ax.plot3D(X.numpy(), Y.numpy(), Z.numpy(), 'gray')
  return ax

#def surface(X, Y, Z, rstride=1, cstride=1, ax=None):
#  ax = plt.axes(projection='3d') if ax is None else ax
#  ax.plot_surface(X.numpy(), Y.numpy(), Z.numpy(), rstride=rstride, cstride=cstride)
#  return ax