tf_initializer_bug_report.py

# UNCOMMENT LINES BELOW TO MAKE THINGS WORK AGAIN
# from tensorflow.python.ops import variables
# def passthrough(obj, value): return value
# try:
#   variables.Variable._build_initializer_expr=passthrough
# except: # older versions of TF don't have this
#   pass





















#!/usr/bin/env python
from collections import OrderedDict
from scipy import linalg
from tensorflow.contrib import graph_editor as ge
import contextlib
import inspect
import math
import numpy as np
import os
import socket
import sys
import tensorflow as tf
import time
import traceback

# shortcuts to refer to util module, this lets move external code into
# this module unmodified
util = sys.modules[__name__]   
u = util


default_dtype = tf.float32
USE_MKL_SVD=False                   # Tensorflow vs 

DUMP_BAD_SVD=False                 # when SVD fails, dump matrix to temp



# TODO: speed-up tests by reusing session

args = None  # TODO: replace with object that crashes on access
def set_global_args(local_args):
  global args
  assert args is None
  args = local_args

def concat_blocks(blocks, validate_dims=True):
  """Takes 2d grid of blocks representing matrices and concatenates to single
  matrix (aka ArrayFlatten)"""

  if validate_dims:
    col_dims = np.array([[int(b.shape[1]) for b in row] for row in blocks])
    col_sums = col_dims.sum(1)
    assert (col_sums[0] == col_sums).all()
    row_dims = np.array([[int(b.shape[0]) for b in row] for row in blocks])
    row_sums = row_dims.sum(0)
    assert (row_sums[0] == row_sums).all()
  
  block_rows = [tf.concat(row, axis=1) for row in blocks]
  return tf.concat(block_rows, axis=0)

def concat_blocks_test():
  blocks = [[tf.constant([[1]]), tf.constant([[1,2]])],
            [tf.transpose(tf.constant([[1,2]])), tf.constant([[1,2],[3,4]])]]
  result = concat_blocks(blocks)
  sess = tf.Session()
  result0 = sess.run(result)
  check_equal(result0, [[1, 1, 2], [1, 1, 2], [2, 3, 4]])


def partition_matrix_evenly(mat, splits):
  """Breaks matrix into 2d grid of equal size."""
  assert int(mat.shape[0])%splits==0
  assert int(mat.shape[1])%splits==0
  
  row_chunks = tf.split(mat, splits, axis=0)
  col_chunks = [tf.split(chunk, splits, axis=1) for chunk in row_chunks]
  return col_chunks

def partition_matrix_evenly_test():
  a = tf.reshape([1,2,3,4], (2,2))
  blocks = partition_matrix_evenly(a, 2)
  a2 = concat_blocks(blocks)
  sess = tf.Session()
  check_equal(sess.run(a2), sess.run(a))

# inverse of concat blocks
def partition_matrix(mat, sizes):
  pass

def partition_matrix_test():
  pass


  # TODO: add name property
def pseudo_inverse(mat, eps=1e-10):
  """Computes pseudo-inverse of mat, treating eigenvalues below eps as 0."""
  
  s, u, v = tf.svd(mat)
  eps = 1e-10   # zero threshold for eigenvalues
  si = tf.where(tf.less(s, eps), s, 1./s)
  return u @ tf.diag(si) @ tf.transpose(v)

def symsqrt(mat, eps=1e-7):
  """Symmetric square root."""
  s, u, v = tf.svd(mat)
  # sqrt is unstable around 0, just use 0 in such case
  print("Warning, cutting off at eps")
  si = tf.where(tf.less(s, eps), s, tf.sqrt(s))
  return u @ tf.diag(si) @ tf.transpose(v)

def pseudo_inverse_sqrt(mat, eps=1e-7):
  """half pseduo-inverse"""
  s, u, v = tf.svd(mat)
  # zero threshold for eigenvalues
  si = tf.where(tf.less(s, eps), s, 1./tf.sqrt(s))
  return u @ tf.diag(si) @ tf.transpose(v)

def pseudo_inverse_sqrt2(svd, eps=1e-7):
  """half pseduo-inverse, accepting existing values"""
  # zero threshold for eigenvalues
  if svd.__class__.__name__=='SvdTuple':
    (s, u, v) = (svd.s, svd.u, svd.v)
  elif svd.__class__.__name__=='SvdWrapper':
    (s, u, v) = (svd.s, svd.u, svd.v)
  else:
    assert False, "Unknown type"
  si = tf.where(tf.less(s, eps), s, 1./tf.sqrt(s))
  return u @ tf.diag(si) @ tf.transpose(v)

def pseudo_inverse2(svd, eps=1e-7):
  """pseudo-inverse, accepting existing values"""
  # use float32 machine precision as cut-off (works for MKL)
  # https://www.wolframcloud.com/objects/927b2aa5-de9c-46f5-89fe-c4a58aa4c04b
  if svd.__class__.__name__=='SvdTuple':
    (s, u, v) = (svd.s, svd.u, svd.v)
  elif svd.__class__.__name__=='SvdWrapper':
    (s, u, v) = (svd.s, svd.u, svd.v)
  else:
    assert False, "Unknown type"
  max_eigen = tf.reduce_max(s)
  si = tf.where(s/max_eigen<eps, 0.*s, 1./s)
  return u @ tf.diag(si) @ tf.transpose(v)

def pseudo_inverse_stable(svd, eps=1e-7):
  """pseudo-inverse, accepting existing values"""
  # use float32 machine precision as cut-off (works for MKL)
  # https://www.wolframcloud.com/objects/927b2aa5-de9c-46f5-89fe-c4a58aa4c04b
  if svd.__class__.__name__=='SvdTuple':
    (s, u, v) = (svd.s, svd.u, svd.v)
  elif svd.__class__.__name__=='SvdWrapper':
    (s, u, v) = (svd.s, svd.u, svd.v)
  else:
    assert False, "Unknown type"
  max_eigen = tf.reduce_max(s)
  si = tf.where(s/max_eigen<eps, 0.*s, tf.pow(s, -0.9))
  return u @ tf.diag(si) @ tf.transpose(v)

# todo: rename l to L
def regularized_inverse(mat, l=0.1):
  return tf.matrix_inverse(mat + l*Identity(int(mat.shape[0])))

# TODO: this gives biased result when I use identity
def regularized_inverse2(svd, L=1e-3):
  """Regularized inverse, working from SVD"""
  if svd.__class__.__name__=='SvdTuple' or svd.__class__.__name__=='SvdWrapper':
    (s, u, v) = (svd.s, svd.u, svd.v)
  else:
    assert False, "Unknown type"
  max_eigen = tf.reduce_max(s)
  #  max_eigen = tf.Print(max_eigen, [max_eigen], "max_eigen")
  #si = 1/(s + L*tf.ones_like(s)/max_eigen)
  si = 1/(s+L*tf.ones_like(s))
  return u @ tf.diag(si) @ tf.transpose(v)

def regularized_inverse3(svd, L=1e-3):
  """Unbiased version of regularized_inverse2"""
  if svd.__class__.__name__=='SvdTuple' or svd.__class__.__name__=='SvdWrapper':
    (s, u, v) = (svd.s, svd.u, svd.v)
  else:
    assert False, "Unknown type"

  if L.__class__.__name__=='Var':
    L = L.var
    
  max_eigen = tf.reduce_max(s)
  #  max_eigen = tf.Print(max_eigen, [max_eigen], "max_eigen")
  #si = 1/(s + L*tf.ones_like(s)/max_eigen)
  si = (1+L*tf.ones_like(s))/(s+L*tf.ones_like(s))
  return u @ tf.diag(si) @ tf.transpose(v)

def regularized_inverse4(svd, L=1e-3):
  """Uses relative norm"""
  if svd.__class__.__name__=='SvdTuple' or svd.__class__.__name__=='SvdWrapper':
    (s, u, v) = (svd.s, svd.u, svd.v)
  else:
    assert False, "Unknown type"

  if L.__class__.__name__=='Var':
    L = L.var
    
  max_eigen = tf.reduce_max(s)
  L = L/max_eigen
  si = (1+L*tf.ones_like(s))/(s+L*tf.ones_like(s))
  #  si = tf.ones_like(s)
  return u @ tf.diag(si) @ tf.transpose(v)

def pseudo_inverse_scipy(tensor):
    dtype = tensor.dtype
    print(linalg.pinv, tensor, dtype)
    result = tf.py_func(linalg.pinv, [tensor],
                        [dtype])[0]
    result.set_shape(tensor.shape)
    return result
  
  

def Identity(n, dtype=None, name=None):
  """Identity matrix of size n."""
  if hasattr(n, "shape"):  # got a Tensor
    nn = fix_shape(n.shape)
    assert nn[0] == nn[1]
    n = nn[0]
  if not dtype:
    dtype = default_dtype
  return tf.diag(tf.ones((n,), dtype=dtype), name=name)

def ones(n, dtype=None, name=None):
  if not dtype:
    dtype = default_dtype
  return tf.ones((n,), dtype=dtype, name=name)

# partitions numpy array into sublists of given sizes
def partition_list_np(vec, sizes):
  assert np.sum(sizes) == len(vec)
  splits = []
  current_idx = 0
  for i in range(len(sizes)):
    splits.append(vec[current_idx: current_idx+sizes[i]])
    current_idx += sizes[i]
  assert current_idx == len(vec)
  return splits

def chunks(l, n):
  """Yield successive n-sized chunks from l."""
  for i in range(0, len(l), n):
    yield l[i:i + n]

def partition_list(l, sizes):
  """Partition l into sublists of given sizes."""
  assert len(l.shape) == 1
  assert np.sum(sizes) == l.shape[0]
  splits = []
  current_idx = 0
  for i in range(len(sizes)):
    splits.append(l[current_idx: current_idx+sizes[i]])
    current_idx += sizes[i]
  return splits

def partition_list_test():
  vec = tf.constant([1,2,3,4,5])
  sess = tf.Session()
  result = sess.run(partition_list(vec, [3, 2]))
  check_equal(result[0], [1,2,3])
  assert (result[1] == [4,5]).all()


def v2c(vec):
  """Convert vector to column matrix."""
  assert len(vec.shape) == 1
  return tf.expand_dims(vec, 1)

def v2c_np(vec):
  """Convert vector to column matrix."""
  assert len(vec.shape) == 1
  return np.expand_dims(vec, 1)

def v2r(vec):
  """Convert vector into row matrix."""
  assert len(vec.shape) == 1
  return tf.expand_dims(vec, 0)
  
def c2v(col):
  """Convert vector into row matrix."""
  assert len(col.shape) == 2
  assert col.shape[1] == 1
  return tf.reshape(col, [-1])


def unvectorize_np(vec, rows):
  """Turn vectorized version of tensor into original matrix with given
  number of rows."""
  assert len(vec)%rows==0
  cols = len(vec)//rows;
  return np.array(np.split(vec, cols)).T

def unvec(vec, rows):
  """Turn vectorized version of tensor into original matrix with given
  number of rows."""
  assert len(vec.shape) == 1
  assert vec.shape[0]%rows == 0
  cols = int(vec.shape[0]//rows)
  return tf.transpose(tf.reshape(vec, (cols, -1)))
#  cols = [v2r(v) for v in tf.split(vec, cols)]
#  return tf.transpose(tf.concat(cols, 0))

def unvec_test():
  vec = tf.constant([1,2,3,4,5,6])
  sess = tf.Session()
  result = sess.run(unvec(vec, 2))
  assert (result==[[1,3,5],[2,4,6]]).all()

def vectorize_np(mat):
  return mat.reshape((-1, 1), order="F")

def vec(mat):
  """Vectorize matrix."""
  return tf.reshape(tf.transpose(mat), [-1,1])

def vec_test():
  mat = tf.constant([[1, 3, 5], [2, 4, 6]])
  sess = tf.Session()
  check_equal(sess.run(c2v(vec(mat))), [1,2,3,4,5,6])


def Kmat(rows, cols):
  """Commutation matrix. Kmat(a,b).vec(M) takes vec of a,b matrix M to vec of
  its transpose."""
  input_mat = np.reshape(np.arange(rows*cols),[rows,-1]).astype(np.int32)
  output_mat = input_mat.T
    
  input_vec = vectorize_np(input_mat)
  output_vec = vectorize_np(output_mat)
    
  K = np.zeros((rows*cols, rows*cols), dtype=np.int32)
  for output_idx in range(rows*cols):
    for input_idx in range(rows*cols):
      K[output_idx, input_idx] = (output_vec[output_idx] == input_vec[input_idx])
  return K

def Kmat_test():
  check_equal(Kmat(3,2),
              [[1, 0, 0, 0, 0, 0],
               [0, 0, 0, 1, 0, 0],
               [0, 1, 0, 0, 0, 0],
               [0, 0, 0, 0, 1, 0],
               [0, 0, 1, 0, 0, 0],
               [0, 0, 0, 0, 0, 1]])

  check_equal(Kmat(2,3),
              [[1, 0, 0, 0, 0, 0],
               [0, 0, 1, 0, 0, 0],
               [0, 0, 0, 0, 1, 0],
               [0, 1, 0, 0, 0, 0],
               [0, 0, 0, 1, 0, 0],
               [0, 0, 0, 0, 0, 1]])

# turns flattened representation into list of matrices with given matrix
# sizes
def unflatten_np(Wf, fs):
  if len(Wf.shape)==2 and Wf.shape[1] == 1:  # treat col mats as vectors
    Wf = Wf.reshape(-1)
    
  dims = [(fs[i+1],fs[i]) for i in range(len(fs)-1)]
  sizes = [s[0]*s[1] for s in dims]
  assert np.sum(sizes)==len(Wf)
  Wsf = partition_list_np(Wf, sizes)
  Ws = [unvectorize_np(Wsf[i], dims[i][0]) for i in range(len(sizes))]
  return Ws

def flatten_np(Ws):
  return np.concatenate([np.reshape(vectorize_np(W),(-1,)) for W in Ws],
                          axis=0)
def flatten_np_test():
  vec = np.asarray(range(1, 11))
  fs = [2,2,2,1]
  result = unflatten_np(vec, fs)
  result2 = flatten_np(result)
  check_equal(vec, result2)

def unflatten(Wf, fs):
  """Turn flattened Tensor into list of rank-2 tensors with given sizes."""
  
  Wf_shape = fix_shape(Wf.shape)
  if len(Wf_shape)==2 and Wf_shape[1] == 1:  # treat col mats as vectors
    Wf = tf.reshape(Wf, [-1])
  dims = [(fs[i+1],fs[i]) for i in range(len(fs)-1)]
  sizes = [s[0]*s[1] for s in dims]
  assert len(Wf.shape) == 1
  assert np.sum(sizes)==Wf.shape[0]
  Wsf = partition_list(Wf, sizes)
  Ws = [unvec(Wsf[i], dims[i][0]) for i in range(len(sizes))]
  return Ws

def unflatten_test():
  vec = tf.constant(list(range(1, 11)))
  sess = tf.Session()
  fs = [2,2,2,1]
  result = sess.run(unflatten(vec, fs))
  check_equal(result[0], [[1,3],[2,4]])
  check_equal(result[1], [[5,7],[6,8]])
  check_equal(result[2], [[9, 10]])

def flatten(Ws):
  """Inverse of unflatten."""
  return tf.concat([tf.reshape(vec(W),(-1,)) for W in Ws], axis=0)

def flatten_test():
  vec = tf.constant(list(range(1, 11)))
  sess = tf.Session()
  fs = [2,2,2,1]
  result = unflatten(vec, fs)
  result2 = flatten(result)
  check_equal(sess.run(vec), sess.run(result2))

def check_close(a0, b0):
  return check_equal(a0, b0, rtol=1e-5, atol=1e-9)
  
def check_equal(a0, b0, rtol=1e-9, atol=1e-12):
  """Helper function to check that two vectors are equal. If inputs are Tensors
  will evaluate them in default session."""


  a = a0.eval() if hasattr(a0, "eval") else a0
  b = b0.eval() if hasattr(b0, "eval") else b0

  check_passed = True
  try:
    np.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
  except Exception as e:
    check_passed = False
    print("Error" + "-"*60)
    for line in traceback.format_stack():
      print(line.strip())
        
    exc_type, exc_value, exc_traceback = sys.exc_info()
    print("*** print_tb:")
    traceback.print_tb(exc_traceback, limit=10, file=sys.stdout)
    efmt = traceback.format_exc()
    print(efmt)
    #    import pdb; pdb.set_trace()

  return check_passed

# TensorShape([Dimension(2), Dimension(10)]) => (2, 10)
def fix_shape(tf_shape):
  return tuple(int(dim) for dim in tf_shape)

def kronecker_cols(a, b):
  """Treats rank-1 vectors a, b as columns, returns Kronecker product a x b."""
  
  assert len(a.get_shape())==1, "Input a must be rank-1, got shape %s" %(a.get_shape(),)
  assert len(b.get_shape())==1, "Input b must be rank-1, got shape %s"%(a.get_shape(),)
  segments = []
  for i in range(a.get_shape()[0]):
    segments.append(a[i]*b)
  result_vec = tf.concat(segments, axis=0)
  result_col = tf.expand_dims(result_vec, 1)
  return result_col

def kronecker_cols_test():
  a = tf.constant([1,2])
  b = tf.constant([3,4])
  c = tf.transpose(tf.constant([[3,4,6,8]]))
  sess = tf.Session()
  assert sess.run(tf.equal(kronecker_cols(a, b), c)).all()


def kronecker(A, B, do_shape_inference=True):
  """Kronecker product of A,B.
  turn_off_shape_inference: if True, makes 10x10 kron go 2.4 sec -> 0.9 sec
  """

  Arows, Acols = fix_shape(A.shape)
  Brows, Bcols = fix_shape(B.shape)
  Crows, Ccols = Arows*Brows, Acols*Bcols
  
  temp = tf.reshape(A, [-1, 1, 1])*tf.expand_dims(B, 0)
  Bshape = tf.constant((Brows, Bcols))

  # turn off shape inference
  if not do_shape_inference:
    disable_shape_inference()

  # [1, n, m] => [n, m]
  slices = [tf.reshape(s, Bshape) for s in tf.split(temp, Crows)]
  
  #  import pdb; pdb.set_trace()
  grid = list(chunks(slices, Acols))
  assert len(grid) == Arows
  result = concat_blocks(grid, validate_dims=do_shape_inference)

  if not do_shape_inference:
    enable_shape_inference()
    result.set_shape((Arows*Brows, Acols*Bcols))
    
  return result

kr = kronecker

def kronecker_test():
  A0 = [[1,2],[3,4]]
  B0 = [[6,7],[8,9]]
  A = tf.constant(A0)
  B = tf.constant(B0)
  C = kronecker(A, B)
  sess = tf.Session()
  C0 = sess.run(C)
  Ct = [[6, 7, 12, 14], [8, 9, 16, 18], [18, 21, 24, 28], [24, 27, 32, 36]]
  Cnp = np.kron(A0, B0)
  check_equal(C0, Ct)
  check_equal(C0, Cnp)


def col(A,i):
  """Extracts i'th column of matrix A"""
  assert len(A.get_shape())==2
  assert i>=0 and i < A.get_shape()[1]
  return tf.expand_dims(A[:,i], 1)


def khatri_rao(A, B):
  Arows, Acols = fix_shape(A.shape)
  Brows, Bcols = fix_shape(B.shape)
  assert Acols==Bcols
  return tf.reshape(tf.einsum("ik,jk->ijk", A, B), (Arows*Brows, Acols))


def khatri_rao_test():
  A = tf.constant([[1, 2], [3, 4]])
  B = tf.constant([[5, 6], [7, 8]])
  C = tf.constant([[5,12], [7,16], [15,24], [21,32]])
  sess = tf.Session()
  assert sess.run(tf.equal(khatri_rao(A, B), C)).all()

  
def relu_mask(a, dtype=default_dtype):
  """Produces mask of 1s for positive values and 0s for negative values."""
  from tensorflow.python.ops import gen_nn_ops
  ones = tf.ones(a.get_shape(), dtype=dtype)
  return gen_nn_ops._relu_grad(ones, a)

def relu_mask_test():
  a = tf.constant([-1,0,1,2], dtype=default_dtype)
  sess = tf.Session()
  check_equal(sess.run(relu_mask(a)), [0,0,1,1])

def assert_rectangular(blocks):
  lengths = np.array([len(row) for row in blocks])
  assert (lengths==lengths[0]).all()
  
def empty_grid(rows, cols):
  """Create empty list of lists of rows-by-cols shape."""
  result = []
  for i in range(rows):
    result.append([None]*cols)
  return result

def block_diagonal_inverse(blocks):
  """Invert diagonal blocks, leave remaining unchanged."""
  
  assert_rectangular(blocks)
  num_rows = len(blocks)
  num_cols = len(blocks[0])

  result = empty_grid(num_rows, num_cols)
  dtype = blocks[0][0].dtype   # TODO: assert same dtype
  
  for i in range(len(blocks)):
    for j in range(len(blocks[0])):
      block = blocks[i][j]
      if i == j:
        result[i][j] = pseudo_inverse(block)
      else:
        result[i][j] = tf.zeros(shape=block.get_shape(),
                                dtype=dtype)
  return result
        
def block_diagonal_inverse_sqrt(blocks):
  assert_rectangular(blocks)
  num_rows = len(blocks)
  num_cols = len(blocks[0])

  result = empty_grid(num_rows, num_cols)
  dtype = blocks[0][0].dtype   # TODO: assert same dtype
  
  for i in range(len(blocks)):
    for j in range(len(blocks[0])):
      block = blocks[i][j]
      if i == j:
        result[i][j] = pseudo_inverse_sqrt(block)
      else:
        result[i][j] = tf.zeros(shape=block.get_shape(),
                                dtype=dtype)
  return result


def block_diagonal_inverse_test():
  sess = tf.Session()
  blocks = [[2*Identity(3), tf.ones((3, 1))],
              [tf.ones((1,3)), 2*Identity(1)]]
  new_blocks = block_diagonal_inverse(blocks)
  actual = concat_blocks(new_blocks)
  expected = 0.5*Identity(4)
  check_equal(sess.run(actual), sess.run(expected))

  
def t(x):
  return tf.transpose(x)

  
# Time tracking functions
global_time_list = []
global_last_time = 0
def reset_time():
  global global_time_list, global_last_time
  global_time_list = []
  global_last_time = time.perf_counter()
  
def record_time():
  global global_last_time, global_time_list
  new_time = time.perf_counter()
  global_time_list.append(new_time - global_last_time)
  global_last_time = time.perf_counter()

def last_time():
  global global_last_time, global_time_list
  if global_time_list:
    return 1000*global_time_list[-1]
  else:
    return 0

def summarize_time(time_list=None):
  if time_list is None:
    time_list = global_time_list

  if time_list[0]>3600*10:   # large first interval means no relative ts
    del time_list[0]
    
  time_list = 1000*np.array(time_list)  # get seconds, convert to ms
  if len(time_list)>0:
    min = np.min(time_list)
    median = np.median(time_list)
    formatted = ["%.2f"%(d,) for d in time_list[:10]]
    print("Times: min: %.2f, median: %.2f, mean: %.2f"%(min, median, np.mean(time_list)))
    #    print("Times: min: %.2f, median: %.2f, mean: %.2f"%(min, median,",".join(formatted)))
  else:
    print("Times: <empty>")
    
def summarize_graph(g=None):
  if not g:
    g = tf.get_default_graph()
  print("Graph: %d ops, %d MBs"%(len(g.get_operations()),
                                 len(str(g.as_graph_def()))/10**6))

from tensorflow.python.framework import ops
original_shape_func = ops.set_shapes_for_outputs
def disable_shape_inference():
  ops.set_shapes_for_outputs = lambda _: _
  
def enable_shape_inference():
  ops.set_shapes_for_outputs = original_shape_func

# work-around for graph_editor.copy_with_input_replacements scaling
# quadratically with size of the graph
from tensorflow.contrib.graph_editor import transform
original_assign_renamed_collections_handler = transform.assign_renamed_collections_handler
def dummy_collections_handler(info, elem, elem_): pass
def disable_collections_handler():
  transform.assign_renamed_collections_handler = dummy_collections_handler
def enable_collections_handler():
  transform.assign_renamed_collections_handler = original_assign_renamed_collections_handler


def dump_with_prompt(result, fname, no_prefix=False):
  """Helper function to ask for confirmation before overwriting."""
  location = os.getcwd()+"/data/"+fname  # TODO: factor out locations logic
  if os.path.exists(location):
    answer = input("%s exists, overwrite? (Y/n) "%(location,))
    if not answer:
      answer = "y"
    if answer.lower() != "y":
      print("skipping")
    else:
      u.dump(result, fname, no_prefix)
  else:
    u.dump(result, fname, no_prefix)
    

def dump(result, fname, no_prefix=False):
  """Save result to file."""
  result = result.eval() if hasattr(result, "eval") else result
  result = np.asarray(result)
  if result.shape == ():   # savetxt has problems with scalars
    result = np.expand_dims(result, 0)
  if no_prefix:
    location = os.getcwd()+"/"+fname
  else:
    location = os.getcwd()+"/data/"+fname
  # special handling for integer datatypes
  if (
      result.dtype == np.uint8 or result.dtype == np.int8 or
      result.dtype == np.uint16 or result.dtype == np.int16 or
      result.dtype == np.uint32 or result.dtype == np.int32 or
      result.dtype == np.uint64 or result.dtype == np.int64
  ):
    np.savetxt(location, result, fmt="%d", delimiter=',')
  else:
    np.savetxt(location, result, delimiter=',')
  print(location)

def dump32(result, fname):
  """Efficient dumping of float32 vals"""
  result = result.eval() if hasattr(result, "eval") else result
  result = np.asarray(result)
  location = os.getcwd()+"/data/"+fname
  assert is_numeric(result)
#  print(location)
  return result.astype('float32').tofile(location)


def frobenius_np(a):
  return np.sqrt(np.sum(np.square(a)))

def nan_check(result):
  result = result.eval() if hasattr(result, "eval") else result
  result = np.asarray(result)
  print("result any NaNs: %s"% (np.isnan(result).any(),))


def L2(t):
  """Squared L2 norm of t."""
  if t.__class__.__name__=='Grads':
    t = t.f
  else:
    assert (t.__class__.__name__ == 'Tensor' or
            t.__class__.__name__.endswith('Variable'))
  return tf.reduce_sum(tf.square(t))


global_timeit_dict = OrderedDict()
class timeit:
  """Decorator to measure length of time spent in the block in millis and log
  it to TensorBoard."""
  
  def __init__(self, tag=""):
    self.tag = tag
    
  def __enter__(self):
    self.start = time.perf_counter()
    return self
  
  def __exit__(self, *args):
    self.end = time.perf_counter()
    interval_ms = 1000*(self.end - self.start)
    global_timeit_dict.setdefault(self.tag, []).append(interval_ms)
    print("    %s Elapsed: %.2f ms"%(self.tag, interval_ms))
    logger = u.get_last_logger(skip_existence_check=True)
    if logger:
      newtag = 'time/'+self.tag
      # since tensorboard doesn't allow hierarchical tags, merge init times
      if newtag.startswith('time/init'):
        newtag = newtag.replace('time/init', 'timeinit')
      logger(newtag, interval_ms)


global_record_dict = OrderedDict()
def record(tag, stat):
    global global_record_dict
    global_record_dict.setdefault(tag, []).append(stat)


def timeit_summarize():
  global global_timeit_dict
  pass

# graph traversal
# computation flows from parents to children
# to find path from target to dependency, do
# nx.shortest_path(gg, dependency, target)
def parents(op): return set(input.op for input in op.inputs)
def children(op): return set(op for out in op.outputs for op in out.consumers())
def dict_graph():
  """Creates dictionary {node: {child1, child2, ..},..} for current
  TensorFlow graph. Result is compatible with networkx/toposort"""

  ops = tf.get_default_graph().get_operations()
  return {op: children(op) for op in ops}
def nx_graph():
  return nx.DiGraph(dict_graph())

def shortest_path(dep, target):
  if hasattr(dep, "op"):
    dep = dep.op
  if hasattr(target, "op"):
    target = target.op
  return nx.shortest_path(nx_graph(), dep, target)

def list_or_tuple(k):
  return isinstance(k, list) or isinstance(k, tuple)

def is_numeric(ndarray):
  ndarray = np.asarray(ndarray)
  return np.issubdtype(ndarray.dtype, np.number)

class VarInfo:
  """Encapsulate variable info."""
  def __init__(self, setter, p):
    self.setter = setter
    self.p = p

class SvdTuple:
  """Object to store svd tuple.
  Create as SvdTuple((s,u,v)) or SvdTuple(s, u, v).
  """
  def __init__(self, suvi, *args):
    if list_or_tuple(suvi):
      if len(suvi) == 3:
        s, u, v = suvi
        inv = Identity(s.shape[0])
      else:
        s, u, v, inv = suvi
    else:
      s = suvi
      u = args[0]
      v = args[1]
      if len(args)>2:
        inv = args[2]
      else:
        inv = Identity(s.shape[0])
    self.s = s
    self.u = u
    self.v = v
    self.inv = inv


class SvdWrapper:
  """Encapsulates variables needed to perform SVD of a TensorFlow target.
  Initialize: wrapper = SvdWrapper(tensorflow_var)
  Trigger SVD: wrapper.update_tf() or wrapper.update_scipy()
  Access result as TF vars: wrapper.s, wrapper.u, wrapper.v
  """
  
  def __init__(self, target, name, do_inverses=False):
    self.name = name
    self.target = target
    self.do_inverses = do_inverses
    self.tf_svd = SvdTuple(tf.svd(target))
    self.update_counter = 0

    self.init = SvdTuple(
      ones(target.shape[0], name=name+"_s_init"),
      Identity(target.shape[0], name=name+"_u_init"),
      Identity(target.shape[0], name=name+"_v_init"),
      Identity(target.shape[0], name=name+"_inv_init"),
    )

    assert self.tf_svd.s.shape == self.init.s.shape
    assert self.tf_svd.u.shape == self.init.u.shape
    assert self.tf_svd.v.shape == self.init.v.shape
    #    assert self.tf_svd.inv.shape == self.init.inv.shape

    self.cached = SvdTuple(
      tf.Variable(self.init.s, name=name+"_s"),
      tf.Variable(self.init.u, name=name+"_u"),
      tf.Variable(self.init.v, name=name+"_v"),
      tf.Variable(self.init.inv, name=name+"_inv"),
    )

    self.s = self.cached.s
    self.u = self.cached.u
    self.v = self.cached.v
    self.inv = self.cached.inv
    
    self.holder = SvdTuple(
      tf.placeholder(default_dtype, shape=self.cached.s.shape, name=name+"_s_holder"),
      tf.placeholder(default_dtype, shape=self.cached.u.shape, name=name+"_u_holder"),
      tf.placeholder(default_dtype, shape=self.cached.v.shape, name=name+"_v_holder"),
      tf.placeholder(default_dtype, shape=self.cached.inv.shape, name=name+"_inv_holder")
    )

    self.update_tf_op = tf.group(
      self.cached.s.assign(self.tf_svd.s),
      self.cached.u.assign(self.tf_svd.u),
      self.cached.v.assign(self.tf_svd.v),
      self.cached.inv.assign(self.tf_svd.inv)
    )

    self.update_external_op = tf.group(
      self.cached.s.assign(self.holder.s),
      self.cached.u.assign(self.holder.u),
      self.cached.v.assign(self.holder.v),
    )

    self.update_externalinv_op = tf.group(
      self.cached.inv.assign(self.holder.inv),
    )


    self.init_ops = (self.s.initializer, self.u.initializer, self.v.initializer,
                     self.inv.initializer)
  

  def update(self):
    if USE_MKL_SVD:
      self.update_scipy()
    else:
      self.update_tf()
    self.update_counter+=1
      
  def update_tf(self):
    sess = tf.get_default_session()
    sess.run(self.update_tf_op)
    
  def update_scipy(self):
    if self.do_inverses:
      return self.update_scipy_inv()
    else:
      return self.update_scipy_svd()

  def update_scipy_inv(self):
    sess = tf.get_default_session()
    target0 = sess.run(self.target)
    inv0 = linalg.inv(target0)
    feed_dict = {self.holder.inv: inv0}
    sess.run(self.update_externalinv_op, feed_dict=feed_dict)
  
  def update_scipy_svd(self):
    sess = tf.get_default_session()
    target0 = sess.run(self.target)
    # A=u.diag(s).v', singular vectors are columns
    # TODO: catch "ValueError: array must not contain infs or NaNs"
    try:
      u0, s0, vt0 = linalg.svd(target0)
      v0 = vt0.T
    except Exception as e:
      print("Got error %s"%(repr(e),))
      if DUMP_BAD_SVD:
        dump32(target0, "badsvd")
      print("gesdd failed, trying gesvd")
      u0, s0, vt0 = linalg.svd(target0, lapack_driver="gesvd")
      v0 = vt0.T
        
    feed_dict = {self.holder.u: u0,
                 self.holder.v: v0,
                 self.holder.s: s0}
    sess.run(self.update_external_op, feed_dict=feed_dict)

def extract_grad(grads_and_vars, var):
  if isinstance(var, str):
    varname = var
  else:
    varname = var.name
  vals = []
  for (grad, var) in grads_and_vars:
    if var.name == varname:
      vals.append(var)
  assert length(vals)==1
  return vals[0]

def intersept_op_creation(op_type_name_to_intercept):
  """Drops into PDB when particular op type is added to graph."""
  from tensorflow.python.framework import op_def_library
  old_apply_op = op_def_library.OpDefLibrary.apply_op
  def my_apply_op(obj, op_type_name, name=None, **keywords):
    print(op_type_name+"-"+str(name))
    if op_type_name == op_type_name_to_intercept:
      import pdb; pdb.set_trace()
    return(old_apply_op(obj, op_type_name, name=name, **keywords))
  op_def_library.OpDefLibrary.apply_op=my_apply_op


global_variables = {}
def get_variable(name, initializer, reuse=True):
  """Lightweight replacement for tf.get_variable() for variables shared within
  a single process. Doesn't need variable scopes."""

  global global_variables
  if name in global_variables and reuse:
    v = global_variables[name]
  else:
    v = tf.Variable(name=name, initial_value=initializer)
    #    print("Creating new variable %s into %s" %(name, v.op.name))
    global_variables[name] = v
  return v


class VarStruct:
  # TODO: refactor to behave more like variable
  """Convenience structure to keep track of variable, its assign op
  and assignment placeholder.

  v = Var(6)
  v.set(5)   # equivalent to sess.run(v.assign_op, feed_dict={pl: 5})
  var.var    # returns underlying variable
  var.val_   # placeholder to assign op
  var.setter # assign op
  var.set(6) # same as sess.run(var.setter, feed_dict={self.val_: val})
  var.initialize()  # sets variable to initial value
  """

  # TODO: add names to placeholder op
  def __init__(self, initial_value, name, dtype=None):

    initial_value = np.array(initial_value)
    assert u.is_numeric(initial_value), "Non-numeric type."
    if not dtype:
      dtype = initial_value.dtype
    else:
      initial_value = initial_value.astype(dtype)
    self.initial_value = initial_value
    self.val_ = tf.placeholder(dtype=initial_value.dtype,
                               shape=initial_value.shape,
                               name=name+"_holder")
    self.var = tf.Variable(initial_value=self.val_, name=name, dtype=dtype)
    assigned_name = self.var.op.name
    if assigned_name != name:
      print("Warning, conflicting variable %s"%(assigned_name,))
    self.setter = self.var.assign(self.val_)

  def set(self, val):
    sess = tf.get_default_session()
    sess.run(self.setter, feed_dict={self.val_: val})

  def initialize(self):
    sess = tf.get_default_session()
    sess.run(self.setter, feed_dict={self.val_: self.val})


global_vars = {}
def get_var(name, initializer, reuse=True):
  """Global get_variable replacement for variables that need to be initialized
  with a large numpy array.
  
  a = tf.get_var([1,2,3])
  a.var   # => gives tf.Variable
  a.val
  """

  global global_vars
  dtype = initializer.dtype
  if name in global_vars and reuse:
    vv = global_vars[name]
    if (np.max(np.abs(vv.initial_value - initializer)))>np.finfo(dtype).eps:
      print("Trying to reinitialize global variable %s with new"
            " value, ignoring new value."%(name,))
  else:
    vv = VarStruct(initial_value=initializer, name=name)
    global_vars[name] = vv
  return vv

def run_all_tests(module):
  all_functions = inspect.getmembers(module, inspect.isfunction)
  for name,func in all_functions:
    if name.endswith("_test"):
      print("Testing "+name)
      with timeit():
        func()
  print(module.__name__+" tests passed.")

@contextlib.contextmanager
def capture_ops():
  """Decorator to capture ops created in the block.
  with capture_ops() as ops:
    # create some ops
  print(ops) # => prints ops created.
  """

  micros = int(time.perf_counter()*10**6)
  scope_name = str(micros)
  op_list = []
  with tf.name_scope(scope_name):
    yield op_list

  g = tf.get_default_graph()
  op_list.extend(ge.select_ops(scope_name+"/.*", graph=g))

@contextlib.contextmanager
def capture_vars():
  """Decorator to capture global variables created in the block.
  """
  
  micros = int(time.perf_counter()*10**6)
  scope_name = "capture_vars_"+str(micros)
  op_list = []
  with tf.variable_scope(scope_name):
    yield op_list

  g = tf.get_default_graph()
  for v in tf.global_variables():
    scope = v.name.split('/', 1)[0]
    if scope == scope_name:
      op_list.append(v)

def Print(op):
  return tf.Print(op, [op], op.name)


def get_host_prefix():
  "ie, returns 10 when on 10.cirrascale..."
  return socket.gethostname().split('.',1)[0]

def summarize_difference(source, target):
  source = np.asarray(source)
  machine_epsilon = np.finfo(source.dtype).eps
  #  abs_diff = np.linalg.norm(np.asarray(source)-target, ord=np.inf)
  abs_diff = abs(np.asarray(source)-target)
  rel_diff = abs_diff/abs(source)/machine_epsilon
  print("abs diff: %f, rel diff: %.1f eps " %(np.max(abs_diff), np.max(rel_diff)))

class BufferedWriter:
  """Class that aggregates multiple writes and flushes periodically."""
  
  def __init__(self, outfn, save_every_secs=60*5):
    self.outfn = outfn
    self.last_save_ts = time.perf_counter()
    self.write_buffer = []
    self.save_every_secs = save_every_secs

  def write(self, line):
    self.write_buffer.append(line)
    if time.perf_counter() - self.last_save_ts > self.save_every_secs:
      self.last_save_ts = time.perf_counter()
      with open(self.outfn, "a") as myfile:
        for line in self.write_buffer:
          myfile.write(line)
      self.write_buffer = []

  def flush():
    with open(outfn, "a") as myfile:
      for line in self.write_buffer:
        myfile.write(line)
    self.write_buffer = []
    
def ossystem(line):
  print(line)
  os.system(line)
  
def setup_experiment_run_directory(run, safe_mode=True):
  # TODO: factor out to use GLOBAL_RUNS_DIRECTORY
  rundir = "runs/%s"%(run,)
  if os.path.exists(rundir):
    if safe_mode and not run=='default':
      answer = input("%s exists, delete? (Y/n) "%(rundir,))
      if not answer:
        answer = "y"
      if answer.lower() != "y":
        print("skipping")
        sys.exit()
    print("Removing %s"%(rundir,))
    ossystem("rm -Rf "+rundir)
  ossystem("mkdir %s"%(rundir,))
  return rundir

########################################
# Tensorboard logging
########################################

# TODO: have global experiment_base that I can use to move logging to
# non-current directory
GLOBAL_RUNS_DIRECTORY='runs'
global_last_logger = None

def get_last_logger(skip_existence_check=False):
  """Returns last logger, if skip_existence_check is set, doesn't
  throw error if logger doesn't exist."""
  global global_last_logger
  if not skip_existence_check:
    assert global_last_logger
  return global_last_logger

class TensorboardLogger:
  """Helper class to log to single tensorboard writer from multiple places.
   logger = u.TensorboardLogger("mnist7")
   logger = u.get_last_logger()  # gets last logger created
   logger('svd_time', 5)  # records "svd_time" stat at 5
   logger.next_step()     # advances step counter
   logger.set_step(5)     # sets step counter to 5
  """
  
  def __init__(self, run, step=0):
    # TODO: do nothing for default run
    
    global global_last_logger
    assert global_last_logger is None
    self.run = run
    #    sess = tf.get_default_session()

    self.summary_writer = tf.summary.FileWriter(GLOBAL_RUNS_DIRECTORY+'/'+run,
                                                graph=tf.get_default_graph())
    self.step = step
    self.summary = tf.Summary()
    global_last_logger = self
    self.last_timestamp = time.perf_counter()

  def __call__(self, *args):
    assert len(args)%2 == 0
    for (tag, value) in chunks(args, 2):
      self.summary.value.add(tag=tag, simple_value=float(value))

  def next_step(self):
    new_timestamp = time.perf_counter()
    interval_ms = 1000*(new_timestamp - self.last_timestamp)
    self.summary.value.add(tag='time/step',
                           simple_value=interval_ms)
    self.last_timestamp = new_timestamp
    self.summary_writer.add_summary(self.summary, self.step)
    self.step+=1
    self.summary = tf.Summary()


def as_int32(v):
  """Convert to int32 dtype."""
  return np.dtype(np.int32).type(v)

def add_dep(from_op, on_op):
  ge.reroute.add_control_inputs(from_op, [on_op])

# Three functions below are replacements for tf default session/default graph
# mechanisms that are global (native ones are thread-local because of thread
# safety issues that have since been fixes (ie, mrry fixed Graph to be thread
# safe for reading)


u = sys.modules[__name__]
util = u

drop_l2 = True               # drop L2 term
drop_sparsity = True         # drop KL term
use_gpu = True
do_line_search = False       # line-search and dump values at each iter

whitening_mode = 4                 # 0 for gradient, 4 for full whitening
whiten_every_n_steps = 1           # how often to whiten
report_frequency = 3               # how often to print loss

num_steps = 20000 if whitening_mode==0 else 20
util.USE_MKL_SVD=True                   # Tensorflow vs MKL SVD

purely_linear = False  # convert sigmoids into linear nonlinearities
use_tikhonov = True    # use Tikhonov reg instead of Moore-Penrose pseudo-inv
Lambda = 1e-3          # magic lambda value from Jimmy Ba for Tikhonov

# adaptive line search
adaptive_step = False     # adjust step length based on predicted decrease
adaptive_step_frequency = 1 # how often to adjust
adaptive_step_burn_in = 0 # let optimization go for a bit before adjusting
local_quadratics = False  # use quadratic approximation to predict loss drop


import tensorflow as tf
#from util import t  # transpose
t = util.t

def W_uniform(s1, s2): # uniform weight init from Ng UFLDL
  r = np.sqrt(6) / np.sqrt(s1 + s2 + 1)
  result = np.random.random(2*s2*s1)*2*r-r
  return result


def main():
  np.random.seed(0)
  tf.set_random_seed(0)
  
  dtype = np.float32
  # 64-bit doesn't help much, search for 64-bit in
  # https://www.wolframcloud.com/objects/5f297f41-30f7-4b1b-972c-cac8d1f8d8e4
  u.default_dtype = dtype
  machine_epsilon = np.finfo(dtype).eps # 1e-7 or 1e-16
  dsize = 10000
  train_images = np.random.random([28*28, dsize])
  patches = train_images[:,:dsize];
  fs = [dsize, 28*28, 196, 28*28]

  # values from deeplearning.stanford.edu/wiki/index.php/UFLDL_Tutorial
  X0=patches
  lambda_=3e-3
  rho=tf.constant(0.1, dtype=dtype)
  beta=3
  W0f = W_uniform(fs[2],fs[3])

  def f(i): return fs[i+1]  # W[i] has shape f[i] x f[i-1]
  dsize = f(-1)
  n = len(fs) - 2

  # helper to create variables with numpy or TF initial value
  init_dict = {}     # {var_placeholder: init_value}
  vard = {}          # {var: util.VarInfo}
  def init_var(val, name, trainable=False, noinit=False):
    if isinstance(val, tf.Tensor):
      collections = [] if noinit else None
      var = tf.Variable(val, name=name, collections=collections)
    else:
      val = np.array(val)
      assert u.is_numeric, "Unknown type"
      holder = tf.placeholder(dtype, shape=val.shape, name=name+"_holder")
      var = tf.Variable(holder, name=name, trainable=trainable)
      init_dict[holder] = val
    var_p = tf.placeholder(var.dtype, var.shape)
    var_setter = var.assign(var_p)
    vard[var] = u.VarInfo(var_setter, var_p)
    return var

  lr = init_var(0.2, "lr")
  if purely_linear:   # need lower LR without sigmoids
    lr = init_var(.02, "lr")
    
  Wf = init_var(W0f, "Wf", True)
  Wf_copy = init_var(W0f, "Wf_copy", True)
  W = u.unflatten(Wf, fs[1:])   # perftodo: this creates transposes
  X = init_var(X0, "X")
  W.insert(0, X)

  def sigmoid(x):
    if not purely_linear:
      return tf.sigmoid(x)
    else:
      return tf.identity(x)
      
  def d_sigmoid(y):
    if not purely_linear:
      return y*(1-y)
    else:
      return 1
    
  def kl(x, y):
    return x * tf.log(x / y) + (1 - x) * tf.log((1 - x) / (1 - y))
  def d_kl(x, y):
    return (1-x)/(1-y) - x/y
  
  # A[i] = activations needed to compute gradient of W[i]
  # A[n+1] = network output
  A = [None]*(n+2)

  # A[0] is just for shape checks, assert fail on run
  # tf.assert always fails because of static assert
  # fail_node = tf.assert_equal(1, 0, message="too huge")
  fail_node = tf.Print(0, [0], "fail, this must never run")
  with tf.control_dependencies([fail_node]):
    A[0] = u.Identity(dsize, dtype=dtype)
  A[1] = W[0]
  for i in range(1, n+1):
    A[i+1] = sigmoid(W[i] @ A[i])
    
  # reconstruction error and sparsity error
  err = (A[3] - A[1])
  rho_hat = tf.reduce_sum(A[2], axis=1, keep_dims=True)/dsize

  # B[i] = backprops needed to compute gradient of W[i]
  # B2[i] = backprops from sampled labels needed for natural gradient
  B = [None]*(n+1)
  B2 = [None]*(n+1)
  B[n] = err*d_sigmoid(A[n+1])
  sampled_labels_live = tf.random_normal((f(n), f(-1)), dtype=dtype, seed=0)
  sampled_labels = init_var(sampled_labels_live, "sampled_labels", noinit=True)
  B2[n] = sampled_labels*d_sigmoid(A[n+1])
  for i in range(n-1, -1, -1):
    backprop = t(W[i+1]) @ B[i+1]
    backprop2 = t(W[i+1]) @ B2[i+1]
    if i == 1 and not drop_sparsity:
      backprop += beta*d_kl(rho, rho_hat)
      backprop2 += beta*d_kl(rho, rho_hat)
    B[i] = backprop*d_sigmoid(A[i+1])
    B2[i] = backprop2*d_sigmoid(A[i+1])

  # dW[i] = gradient of W[i]
  dW = [None]*(n+1)
  pre_dW = [None]*(n+1)  # preconditioned dW
  pre_dW_stable = [None]*(n+1)  # preconditioned stable dW

  cov_A = [None]*(n+1)    # covariance of activations[i]
  cov_B2 = [None]*(n+1)   # covariance of synthetic backprops[i]
  vars_svd_A = [None]*(n+1)
  vars_svd_B2 = [None]*(n+1)
  for i in range(1,n+1):
    cov_A[i] = init_var(A[i]@t(A[i])/dsize, "cov_A%d"%(i,))
    cov_B2[i] = init_var(B2[i]@t(B2[i])/dsize, "cov_B2%d"%(i,))
    vars_svd_A[i] = u.SvdWrapper(cov_A[i],"svd_A_%d"%(i,))
    vars_svd_B2[i] = u.SvdWrapper(cov_B2[i],"svd_B2_%d"%(i,))
    if use_tikhonov:
      whitened_A = u.regularized_inverse2(vars_svd_A[i],L=Lambda) @ A[i]
    else:
      whitened_A = u.pseudo_inverse2(vars_svd_A[i]) @ A[i]
    if use_tikhonov:
      whitened_B2 = u.regularized_inverse2(vars_svd_B2[i],L=Lambda) @ B[i]
    else:
      whitened_B2 = u.pseudo_inverse2(vars_svd_B2[i]) @ B[i]
    whitened_A_stable = u.pseudo_inverse_sqrt2(vars_svd_A[i]) @ A[i]
    whitened_B2_stable = u.pseudo_inverse_sqrt2(vars_svd_B2[i]) @ B[i]
    pre_dW[i] = (whitened_B2 @ t(whitened_A))/dsize
    pre_dW_stable[i] = (whitened_B2_stable @ t(whitened_A_stable))/dsize
    dW[i] = (B[i] @ t(A[i]))/dsize

  # Loss function
  reconstruction = u.L2(err) / (2 * dsize)
  sparsity = beta * tf.reduce_sum(kl(rho, rho_hat))
  L2 = (lambda_ / 2) * (u.L2(W[1]) + u.L2(W[1]))

  loss = reconstruction
  if not drop_l2:
    loss = loss + L2
  if not drop_sparsity:
    loss = loss + sparsity

  grad_live = u.flatten(dW[1:])
  pre_grad_live = u.flatten(pre_dW[1:]) # fisher preconditioned gradient
  pre_grad_stable_live = u.flatten(pre_dW_stable[1:]) # sqrt fisher preconditioned grad
  grad = init_var(grad_live, "grad")
  pre_grad = init_var(pre_grad_live, "pre_grad")
  pre_grad_stable = init_var(pre_grad_stable_live, "pre_grad_stable")

  update_params_op = Wf.assign(Wf-lr*pre_grad).op
  update_params_stable_op = Wf.assign(Wf-lr*pre_grad_stable).op
  save_params_op = Wf_copy.assign(Wf).op
  pre_grad_dot_grad = tf.reduce_sum(pre_grad*grad)
  pre_grad_stable_dot_grad = tf.reduce_sum(pre_grad*grad)
  grad_norm = tf.reduce_sum(grad*grad)
  pre_grad_norm = u.L2(pre_grad)
  pre_grad_stable_norm = u.L2(pre_grad_stable)

  def dump_svd_info(step):
    """Dump singular values and gradient values in those coordinates."""
    for i in range(1, n+1):
      svd = vars_svd_A[i]
      s0, u0, v0 = sess.run([svd.s, svd.u, svd.v])
      util.dump(s0, "A_%d_%d"%(i, step))
      A0 = A[i].eval()
      At0 = v0.T @ A0
      util.dump(A0 @ A0.T, "Acov_%d_%d"%(i, step))
      util.dump(At0 @ At0.T, "Atcov_%d_%d"%(i, step))
      util.dump(s0, "As_%d_%d"%(i, step))

    for i in range(1, n+1):
      svd = vars_svd_B2[i]
      s0, u0, v0 = sess.run([svd.s, svd.u, svd.v])
      util.dump(s0, "B2_%d_%d"%(i, step))
      B0 = B[i].eval()
      Bt0 = v0.T @ B0
      util.dump(B0 @ B0.T, "Bcov_%d_%d"%(i, step))
      util.dump(Bt0 @ Bt0.T, "Btcov_%d_%d"%(i, step))
      util.dump(s0, "Bs_%d_%d"%(i, step))      
    
  def advance_batch():
    sess.run(sampled_labels.initializer)  # new labels for next call

  def update_covariances():
    ops_A = [cov_A[i].initializer for i in range(1, n+1)]
    ops_B2 = [cov_B2[i].initializer for i in range(1, n+1)]
    sess.run(ops_A+ops_B2)

  def update_svds():
    if whitening_mode>1:
      vars_svd_A[2].update()
    if whitening_mode>2:
      vars_svd_B2[2].update()
    if whitening_mode>3:
      vars_svd_B2[1].update()

  def init_svds():
    """Initialize our SVD to identity matrices."""
    ops = []
    for i in range(1, n+1):
      ops.extend(vars_svd_A[i].init_ops)
      ops.extend(vars_svd_B2[i].init_ops)
    sess = tf.get_default_session()
    sess.run(ops)
      
  init_op = tf.global_variables_initializer()
  #  tf.get_default_graph().finalize()
  
  sess = tf.InteractiveSession()
  sess.run(Wf.initializer, feed_dict=init_dict)
  sess.run(X.initializer, feed_dict=init_dict)
  advance_batch()
  update_covariances()
  init_svds()
  sess.run(init_op, feed_dict=init_dict)  # initialize everything else
  
  print("Running training.")
  u.reset_time()

  step_lengths = []     # keep track of learning rates
  losses = []
  ratios = []           # actual loss decrease / expected decrease
  grad_norms = []       
  pre_grad_norms = []   # preconditioned grad norm squared
  pre_grad_stable_norms = [] # sqrt preconditioned grad norms squared
  target_delta_list = []     # predicted decrease linear approximation
  target_delta2_list = []    # predicted decrease quadratic appromation
  actual_delta_list = []      # actual decrease
  
  # adaptive line search parameters
  alpha=0.3   # acceptable fraction of predicted decrease
  beta=0.8    # how much to shrink when violation
  growth_rate=1.05  # how much to grow when too conservative
    
  def update_cov_A(i):
    sess.run(cov_A[i].initializer)
  def update_cov_B2(i):
    sess.run(cov_B2[i].initializer)

  # only update whitening matrix of input activations in the beginning
  if whitening_mode>0:
    vars_svd_A[1].update()

  # compute t(delta).H.delta/2
  def hessian_quadratic(delta):
    #    update_covariances()
    W = u.unflatten(delta, fs[1:])
    W.insert(0, None)
    total = 0
    for l in range(1, n+1):
      decrement = tf.trace(t(W[l])@cov_B2[l]@W[l]@cov_A[l])
      total+=decrement
    return (total/2).eval()
    
  # compute t(delta).H^-1.delta/2
  def hessian_quadratic_inv(delta):
    #    update_covariances()
    W = u.unflatten(delta, fs[1:])
    W.insert(0, None)
    total = 0
    for l in range(1, n+1):
      invB2 = u.pseudo_inverse2(vars_svd_B2[l])
      invA = u.pseudo_inverse2(vars_svd_A[l])
      decrement = tf.trace(t(W[l])@invB2@W[l]@invA)
      total+=decrement
    return (total/2).eval()

  # do line search, dump values as csv
  def line_search(initial_value, direction, step, num_steps):
    saved_val = tf.Variable(Wf)
    sess.run(saved_val.initializer)
    pl = tf.placeholder(dtype, shape=(), name="linesearch_p")
    assign_op = Wf.assign(initial_value - direction*step*pl)
    vals = []
    for i in range(num_steps):
      sess.run(assign_op, feed_dict={pl: i})
      vals.append(loss.eval())
    sess.run(Wf.assign(saved_val)) # restore original value
    return vals
    
  for step in range(num_steps): 
    update_covariances()
    if step % whiten_every_n_steps==0:
      update_svds()

    sess.run(grad.initializer)
    sess.run(pre_grad.initializer)
    
    lr0, loss0 = sess.run([lr, loss])
    save_params_op.run()

    # regular inverse becomes unstable when grad norm exceeds 1
    stabilized_mode = grad_norm.eval()<1

    if stabilized_mode and not use_tikhonov:
      update_params_stable_op.run()
    else:
      update_params_op.run()

    loss1 = loss.eval()
    advance_batch()

    # line search stuff
    target_slope = (-pre_grad_dot_grad.eval() if stabilized_mode else
                    -pre_grad_stable_dot_grad.eval())
    target_delta = lr0*target_slope
    target_delta_list.append(target_delta)

    # second order prediction of target delta
    # TODO: the sign is wrong, debug this
    # https://www.wolframcloud.com/objects/8f287f2f-ceb7-42f7-a599-1c03fda18f28
    if local_quadratics:
      x0 = Wf_copy.eval()
      x_opt = x0-pre_grad.eval()
      # computes t(x)@H^-1 @(x)/2
      y_opt = loss0 - hessian_quadratic_inv(grad)
      # computes t(x)@H @(x)/2
      y_expected = hessian_quadratic(Wf-x_opt)+y_opt
      target_delta2 = y_expected - loss0
      target_delta2_list.append(target_delta2)
      
    
    actual_delta = loss1 - loss0
    actual_slope = actual_delta/lr0
    slope_ratio = actual_slope/target_slope  # between 0 and 1.01
    actual_delta_list.append(actual_delta)

    if do_line_search:
      vals1 = line_search(Wf_copy, pre_grad, lr/100, 40)
      vals2 = line_search(Wf_copy, grad, lr/100, 40)
      u.dump(vals1, "line1-%d"%(i,))
      u.dump(vals2, "line2-%d"%(i,))
      
    losses.append(loss0)
    step_lengths.append(lr0)
    ratios.append(slope_ratio)
    grad_norms.append(grad_norm.eval())
    pre_grad_norms.append(pre_grad_norm.eval())
    pre_grad_stable_norms.append(pre_grad_stable_norm.eval())

    if step % report_frequency == 0:
      print("Step %d loss %.2f, target decrease %.3f, actual decrease, %.3f ratio %.2f grad norm: %.2f pregrad norm: %.2f"%(step, loss0, target_delta, actual_delta, slope_ratio, grad_norm.eval(), pre_grad_norm.eval()))
    
    if adaptive_step_frequency and adaptive_step and step>adaptive_step_burn_in:
      # shrink if wrong prediction, don't shrink if prediction is tiny
      if slope_ratio < alpha and abs(target_delta)>1e-6 and adaptive_step:
        print("%.2f %.2f %.2f"%(loss0, loss1, slope_ratio))
        print("Slope optimality %.2f, shrinking learning rate to %.2f"%(slope_ratio, lr0*beta,))
        sess.run(vard[lr].setter, feed_dict={vard[lr].p: lr0*beta})
        
      # grow learning rate, slope_ratio .99 worked best for gradient
      elif step>0 and i%50 == 0 and slope_ratio>0.90 and adaptive_step:
          print("%.2f %.2f %.2f"%(loss0, loss1, slope_ratio))
          print("Growing learning rate to %.2f"%(lr0*growth_rate))
          sess.run(vard[lr].setter, feed_dict={vard[lr].p:
                                               lr0*growth_rate})

    u.record_time()

  print("Test passed")

if __name__=='__main__':
  main()