# Path: t3f/shapes.py
# def tt_ranks(tt, name='t3f_tt_ranks'):
# def shape(tt, name='t3f_shape'):
# def raw_shape(tt, name='t3f_raw_shape'):
# def batch_size(tt, name='t3f_batch_size'):
# def lazy_tt_ranks(tt, name='t3f_lazy_tt_ranks'):
# def lazy_shape(tt, name='t3f_lazy_shape'):
# def lazy_raw_shape(tt, name='t3f_lazy_raw_shape'):
# def lazy_batch_size(tt, name='t3f_lazy_batch_size'):
# def clean_raw_shape(shape, name='t3f_clean_raw_shape'):
# def is_batch_broadcasting_possible(tt_a, tt_b):
# def squeeze_batch_dim(tt, name='t3f_squeeze_batch_dim'):
# def expand_batch_dim(tt, name='t3f_expand_batch_dim'):
# 
# Path: t3f/decompositions.py
# def to_tt_matrix(mat, shape, max_tt_rank=10, epsilon=None,
#                  name='t3f_to_tt_matrix'):
# def to_tt_tensor(tens, max_tt_rank=10, epsilon=None,
#                  name='t3f_to_tt_tensor'):
# def round(tt, max_tt_rank=None, epsilon=None, name='t3f_round'):
# def _round_tt(tt, max_tt_rank, epsilon):
# def _round_batch_tt(tt, max_tt_rank, epsilon):
# def orthogonalize_tt_cores(tt, left_to_right=True,
#                            name='t3f_orthogonalize_tt_cores'):
# def _orthogonalize_tt_cores_left_to_right(tt):
# def _orthogonalize_batch_tt_cores_left_to_right(tt):
# def _orthogonalize_tt_cores_right_to_left(tt):
# 
# Path: t3f/riemannian.py
# def project_sum(what, where, weights=None):
# def project(what, where):
# def project_matmul(what, where, matrix):
# def pairwise_flat_inner_projected(projected_tt_vectors_1,
#                                   projected_tt_vectors_2):
# def add_n_projected(tt_objects, coef=None):
#   def slice_tt_core(tt_core, left_idx, right_idx):
# def tangent_space_to_deltas(tt, name='t3f_tangent_space_to_deltas'):
# def deltas_to_tangent_space(deltas, tt, left=None, right=None,
#                             name='t3f_deltas_to_tangent_space'):
# 
# Path: t3f/utils.py
# def unravel_index(indices, shape):
# def replace_tf_svd_with_np_svd():
#   def my_svd(tensor, full_matrices=False, compute_uv=True):
# def in_eager_mode():


import tensorflow as tf
from t3f import shapes
from t3f import decompositions
from t3f import riemannian
from t3f import utils



def value_and_grad(f, x):
  """Gradient of the given function w.r.t. x. Works in eager and graph mode."""
  if utils.in_eager_mode():
    with tf.GradientTape() as tape:
      tape.watch(x)
      v = f(x)
    return v, tape.gradient(v, x)
  else:
    v = f(x)
    return v, tf.gradients(v, x)


def _enforce_gauge_conditions(deltas, left):
  """Project deltas that define tangent space vec onto the gauge conditions."""
  proj_deltas = []

# Path: t3f/shapes.py
# def tt_ranks(tt, name='t3f_tt_ranks'):
# def shape(tt, name='t3f_shape'):
# def raw_shape(tt, name='t3f_raw_shape'):
# def batch_size(tt, name='t3f_batch_size'):
# def lazy_tt_ranks(tt, name='t3f_lazy_tt_ranks'):
# def lazy_shape(tt, name='t3f_lazy_shape'):
# def lazy_raw_shape(tt, name='t3f_lazy_raw_shape'):
# def lazy_batch_size(tt, name='t3f_lazy_batch_size'):
# def clean_raw_shape(shape, name='t3f_clean_raw_shape'):
# def is_batch_broadcasting_possible(tt_a, tt_b):
# def squeeze_batch_dim(tt, name='t3f_squeeze_batch_dim'):
# def expand_batch_dim(tt, name='t3f_expand_batch_dim'):
# 
# Path: t3f/decompositions.py
# def to_tt_matrix(mat, shape, max_tt_rank=10, epsilon=None,
#                  name='t3f_to_tt_matrix'):
# def to_tt_tensor(tens, max_tt_rank=10, epsilon=None,
#                  name='t3f_to_tt_tensor'):
# def round(tt, max_tt_rank=None, epsilon=None, name='t3f_round'):
# def _round_tt(tt, max_tt_rank, epsilon):
# def _round_batch_tt(tt, max_tt_rank, epsilon):
# def orthogonalize_tt_cores(tt, left_to_right=True,
#                            name='t3f_orthogonalize_tt_cores'):
# def _orthogonalize_tt_cores_left_to_right(tt):
# def _orthogonalize_batch_tt_cores_left_to_right(tt):
# def _orthogonalize_tt_cores_right_to_left(tt):
# 
# Path: t3f/riemannian.py
# def project_sum(what, where, weights=None):
# def project(what, where):
# def project_matmul(what, where, matrix):
# def pairwise_flat_inner_projected(projected_tt_vectors_1,
#                                   projected_tt_vectors_2):
# def add_n_projected(tt_objects, coef=None):
#   def slice_tt_core(tt_core, left_idx, right_idx):
# def tangent_space_to_deltas(tt, name='t3f_tangent_space_to_deltas'):
# def deltas_to_tangent_space(deltas, tt, left=None, right=None,
#                             name='t3f_deltas_to_tangent_space'):
# 
# Path: t3f/utils.py
# def unravel_index(indices, shape):
# def replace_tf_svd_with_np_svd():
#   def my_svd(tensor, full_matrices=False, compute_uv=True):
# def in_eager_mode():


# Path: t3f/autodiff.py
import tensorflow as tf
from t3f import shapes
from t3f import decompositions
from t3f import riemannian
from t3f import utils




def value_and_grad(f, x):
  """Gradient of the given function w.r.t. x. Works in eager and graph mode."""
  if utils.in_eager_mode():
    with tf.GradientTape() as tape:
      tape.watch(x)
      v = f(x)
    return v, tape.gradient(v, x)
  else:
    v = f(x)
    return v, tf.gradients(v, x)


def _enforce_gauge_conditions(deltas, left):
  """Project deltas that define tangent space vec onto the gauge conditions."""
  proj_deltas = []

  tt_ranks = shapes.lazy_tt_ranks(left)